INNER CODE UNIT · Python

__init__

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

    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)
    model = TestMod()
    model.cuda()
    model.half()
    model.train()
    model = torch.jit.script(model)

View source record →