Commit 7f1c5917 authored by Lawrence McAfee's avatar Lawrence McAfee
Browse files

handle output_tensor == None

parent 8fc5e323

megatron/schedules.py

+2 −0
Original line number Diff line number Diff line
@@ -49,6 +49,8 @@ def deallocate_output_tensor(out): 
    sent to the next pipeline stage. At this point, the output tensor is

    only useful for its '.grad_fn' field, and not its '.data'.

    '''

    if out is None:

        return

    assert isinstance(out, torch.Tensor), \

        "expected Tensor, found %s." % type(out).__name__

    assert out._base is None, \