Skip to content

GitLab

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

handle output_tensor == None

parent 8fc5e323
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 0 deletions
megatron/schedules.py
View file @ 7f1c5917
...@@ -49,6 +49,8 @@ def deallocate_output_tensor(out): ...@@ -49,6 +49,8 @@ def deallocate_output_tensor(out):
sent to the next pipeline stage. At this point, the output tensor is sent to the next pipeline stage. At this point, the output tensor is
only useful for its '.grad_fn' field, and not its '.data'. only useful for its '.grad_fn' field, and not its '.data'.
''' '''
if out is None:
return
assert isinstance(out, torch.Tensor), \ assert isinstance(out, torch.Tensor), \
"expected Tensor, found %s." % type(out).__name__ "expected Tensor, found %s." % type(out).__name__
assert out._base is None, \ assert out._base is None, \
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment