Commit a5bfc013 authored by Neel Kant's avatar Neel Kant
Browse files

Merge branch 'master' into ict-merge

parents 628bf0dd c20f4d48

megatron/mpu/random.py

+3 −1
Original line number Diff line number Diff line
@@ -234,7 +234,9 @@ class CheckpointFunction(torch.autograd.Function): 
        if isinstance(outputs, torch.Tensor):

            outputs = (outputs,)

        torch.autograd.backward(outputs, args)

        return (None,) + tuple(inp.grad for inp in detached_inputs)

        grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp

                      for inp in detached_inputs)

        return (None,) + grads





def checkpoint(function, *args):