INNER CODE UNIT · Python

diff

NVIDIA/DeepLearningExamples · PyTorch/LanguageModeling/BERT/lamb_amp_opt/main.py:21

                diff = (p - master_p.half()).float().abs().mean().item()
                print(f"  {i}th group, {ii}th param diff: {diff}")


class TestMod(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = torch.nn.Sequential(
            torch.nn.Linear(N_FEATURES, N_FEATURES // 2),
            torch.nn.ReLU(),
            torch.nn.Linear(N_FEATURES // 2, 2),
        )
    def forward(self, inputs) :
        return self.layers(inputs)


def main() :
    loss = torch.nn.CrossEntropyLoss(ignore_index=-1)

View source record →