Fix pad_size in pad function

mlcommons · Sep 14, 2023 · ad64fd1 · ad64fd1
1 parent 35c8736
commit ad64fd1
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/algorithmic_efficiency/data_utils.py b/algorithmic_efficiency/data_utils.py
@@ -71,7 +71,7 @@ def _prepare(x):
 def pad(tensor: np.ndarray,
         pad_size: int,
         padding_value: int = 0) -> np.ndarray:
-  if len(tensor) > 1:
+  if tensor.ndim > 1:
     pad_size = (pad_size, *tensor.shape[1:])
   padding = np.full(pad_size, padding_value, dtype=tensor.dtype)
   padded_tensor = np.concatenate((tensor, padding), axis=0)