relax size check in flatten_for_scatter_gather (#40573)

Natalia Gimelshein · facebook-github-bot · commit 74b7f1f7c9f9 · 2020-06-25T11:28:47.000-07:00
Summary: Pull Request resolved: #40573 Per title, to workaround apex sbn bug. Test Plan: Covered by existing tests Reviewed By: blefaudeux Differential Revision: D22236942 fbshipit-source-id: 4e195e35e3e1b309400be53614d5368f6c08801e
diff --git a/torch/lib/c10d/ProcessGroupNCCL.cpp b/torch/lib/c10d/ProcessGroupNCCL.cpp
@@ -624,13 +624,9 @@ std::vector<at::Tensor> flatten_for_scatter_gather(
     }
 
     for (const auto& t : tensor_lists[i]) {
-      if (t.sizes() != other[i].sizes()) {
+      if (t.numel() != other[i].numel()) {
         throw std::runtime_error(
-            "All tensor operands to scatter/gather must have the same size");
-      }
-      if (t.strides() != other[i].strides()) {
-        throw std::runtime_error(
-            "All tensor operands to scatter/gather must have the same layout (strides)");
+            "All tensor operands to scatter/gather must have the same number of elements");
       }
     }
     // Flatten the tensors (from all ranks) into a single big tensor.

Original file line number	Diff line number	Diff line change
`@@ -624,13 +624,9 @@ std::vector<at::Tensor> flatten_for_scatter_gather(`
`624`	`624`	`}`
`625`	`625`
`626`	`626`	`for (const auto& t : tensor_lists[i]) {`
`627`		`- if (t.sizes() != other[i].sizes()) {`
	`627`	`+ if (t.numel() != other[i].numel()) {`
`628`	`628`	`throw std::runtime_error(`
`629`		`- "All tensor operands to scatter/gather must have the same size");`
`630`		`- }`
`631`		`- if (t.strides() != other[i].strides()) {`
`632`		`- throw std::runtime_error(`
`633`		`- "All tensor operands to scatter/gather must have the same layout (strides)");`
	`629`	`+ "All tensor operands to scatter/gather must have the same number of elements");`
`634`	`630`	`}`
`635`	`631`	`}`
`636`	`632`	`// Flatten the tensors (from all ranks) into a single big tensor.`