Skip to content

GitLab

"vscode:/vscode.git/clone" did not exist on "22400011d6a498ecf77797a56dfe13bc94c426ca"
Commit 85876092 authored by Lawrence McAfee's avatar Lawrence McAfee
Browse files

output tensor now allocated directly on gpu.

parent e33d63cc
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
megatron/schedules.py
View file @ 85876092
...@@ -54,7 +54,7 @@ def free_output_tensor(output_tensors): ...@@ -54,7 +54,7 @@ def free_output_tensor(output_tensors):
if isinstance(output_tensors, torch.Tensor): if isinstance(output_tensors, torch.Tensor):
output_tensors = [output_tensors] output_tensors = [output_tensors]
for output_tensor in output_tensors: for output_tensor in output_tensors:
output_tensor.data = torch.FloatTensor([0]).to(output_tensor.data) output_tensor.data = torch.cuda.FloatTensor([0])
def custom_backward(output, grad_output): def custom_backward(output, grad_output):
'''Directly call C++ autograd engine. '''Directly call C++ autograd engine.
......
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