update runtime error message for minibatch (#8243)

pytorch · Oct 30, 2024 · 3efe1eb · 3efe1eb
1 parent dc20b2d
commit 3efe1eb
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/test/spmd/test_xla_sharding.py b/test/spmd/test_xla_sharding.py
@@ -1373,7 +1373,7 @@ def test_data_loader_with_non_batch_size_and_mini_batch(self):
             mesh, ('data', None, None, None), minibatch=True))
     with self.assertRaisesRegex(
         RuntimeError,
-        "When minibatch is configured, batch dimension of the tensor must be divisible by local runtime device count*"
+        "When minibatch is configured, the per-host batch size must be divisible by local runtime device count. Per host input data shape *"
     ):
       data, _ = iter(train_device_loader).__next__()
 

diff --git a/torch_xla/core/xla_model.py b/torch_xla/core/xla_model.py
@@ -1309,10 +1309,9 @@ def convert_fn(tensors):
         if sharding and tensor.dim() > 0 and (tensor.size()[0] %
                                               local_runtime_device_count) != 0:
           raise RuntimeError(
-              "When minibatch is configured, batch dimension of the tensor " +
-              "must be divisible by local runtime device count.input data shape "
-              +
-              f"={tensor.size()}, local_runtime_device_count = {local_runtime_device_count}"
+              "When minibatch is configured, the per-host batch size must be divisible "
+              + "by local runtime device count. Per host input data shape " +
+              f"= {tensor.size()}, local_runtime_device_count = {local_runtime_device_count}"
           )
 
     xtensors = torch_xla._XLAC._xla_tensors_from_aten(tensors, devices,