debugging

mlcommons · Sep 29, 2023 · 6fdc23a · 6fdc23a
1 parent 373fd54
commit 6fdc23a
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/algorithmic_efficiency/data_utils.py b/algorithmic_efficiency/data_utils.py
@@ -36,8 +36,8 @@ def shard_and_maybe_pad_np(
   else:
     pad_to_global_batch_size = False
   remainder_size = current_batch_size % local_device_count
-  if pad_to_global_batch_size:
-  # if remainder_size != 0:
+  # if pad_to_global_batch_size:
+  if remainder_size != 0:
     if global_batch_size is not None:
       pad_size = global_batch_size - current_batch_size
     else:
@@ -60,8 +60,8 @@ def _prepare(x):
       x = x._numpy()  # pylint: disable=protected-access
 
     # Pad if remainder_size != 0 (should only be possible during evaluation).
-    if pad_to_global_batch_size:
-    # if remainder_size != 0: 
+    # if pad_to_global_batch_size:
+    if remainder_size != 0: 
       x = pad(x, pad_size, padding_value=padding_value)
 
     # Reshape (global_batch_size, ...) to