INNER CODE UNIT · Python

forward

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

    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)
    param_optimizer = list(model.named_parameters())
    no_decay = ['bias', 'gamma', 'beta', 'LayerNorm']
    optimizer_grouped_parameters = [
        {'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}]
    grad_scaler = torch.cuda.amp.GradScaler(enabled=True)
    optimizer = FusedLAMBAMP(optimizer_grouped_parameters)

View source record →