Only record grad_fn in C++ Scatter and Gather when required so (#20286)

mrshenli · facebook-github-bot · commit d6815e1e27a8 · 2019-05-08T21:04:21.000-07:00
Summary: C++ `Scatter` and `Gather` always set autograd history for input data tensors regardless whether they require grad. This hits assertion failure in `set_history(Tensor, shared_ptr<Function> grad_fn)` where `grad_fn` cannot be nullptr. After this PR, C++ `Scatter` and `Gather` only record `grad_fn` when required. Pull Request resolved: #20286 Differential Revision: D15266610 Pulled By: mrshenli fbshipit-source-id: 641df0ea36e7c922b5820c8dc3f83e2a050412b5
diff --git a/torch/csrc/autograd/functions/comm.cpp b/torch/csrc/autograd/functions/comm.cpp
@@ -59,7 +59,9 @@ variable_list Scatter::apply(variable_list&& inputs) {
     }
   }
 
-  set_history(variables, grad_fn);
+  if (grad_fn) {
+    set_history(variables, grad_fn);
+  }
 
   return variables;
 }
@@ -120,7 +122,9 @@ variable_list Gather::apply(variable_list&& inputs) {
   const auto destination_index =
       destination_device_.is_cpu() ? -1 : destination_device_.index();
   auto variable = torch::cuda::gather(tensors, dim_, destination_index);
-  set_history(variable, grad_fn);
+  if (grad_fn) {
+    set_history(variable, grad_fn);
+  }
   return {variable};
 }
 

Original file line number	Diff line number	Diff line change
`@@ -59,7 +59,9 @@ variable_list Scatter::apply(variable_list&& inputs) {`
`59`	`59`	`}`
`60`	`60`	`}`
`61`	`61`
`62`		`- set_history(variables, grad_fn);`
	`62`	`+ if (grad_fn) {`
	`63`	`+ set_history(variables, grad_fn);`
	`64`	`+ }`
`63`	`65`
`64`	`66`	`return variables;`
`65`	`67`	`}`
`@@ -120,7 +122,9 @@ variable_list Gather::apply(variable_list&& inputs) {`
`120`	`122`	`const auto destination_index =`
`121`	`123`	`destination_device_.is_cpu() ? -1 : destination_device_.index();`
`122`	`124`	`auto variable = torch::cuda::gather(tensors, dim_, destination_index);`
`123`		`- set_history(variable, grad_fn);`
	`125`	`+ if (grad_fn) {`
	`126`	`+ set_history(variable, grad_fn);`
	`127`	`+ }`
`124`	`128`	`return {variable};`
`125`	`129`	`}`
`126`	`130`