Fix lint

NVIDIA · Nov 28, 2024 · 7ccfccd · 7ccfccd
1 parent 5d60dd5
commit 7ccfccd
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/csrc/multidevice/communication.cpp b/csrc/multidevice/communication.cpp
@@ -326,7 +326,7 @@ c10::intrusive_ptr<c10d::Work> postAllgather(
     at::Tensor input_tensor,
     at::Tensor output_tensor) {
   auto splits =
-      at::tensor_split(output_tensor, communication->team().size(), /*dim=*/0);
+      at::tensor_split(output_tensor, communication->team_size(), /*dim=*/0);
   assertBuffersHaveSameSize({input_tensor}, splits);
 
   // allgather primitive in c10d induces extra buffering time to copy out the

diff --git a/csrc/multidevice/communication.h b/csrc/multidevice/communication.h
@@ -90,6 +90,11 @@ class Communication : public Expr {
     return attribute<Team>(1);
   }
 
+  // A convenience helper so the user doesn't need to convert size_t to int64_t.
+  int64_t team_size() const {
+    return static_cast<int64_t>(team().size());
+  }
+
   DeviceIdxType root() const {
     return attribute<DeviceIdxType>(2);
   }