handle output_tensor == None

MediaPreneur · Jan 13, 2022 · 7f1c591 · 7f1c591
1 parent 8fc5e32
commit 7f1c591
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/megatron/schedules.py b/megatron/schedules.py
@@ -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, \