fix ut

NVIDIA · Aug 29, 2024 · e652a43 · e652a43
1 parent 352c675
commit e652a43
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py
@@ -1437,11 +1437,14 @@ def _generate_random_numbers(n, total_sum):
         if isinstance(block, TorcGroupedLinearWithPadding):
             out = block(inp_hidden_states, m_splits)
         else:
-            padded_inp_hidden_states, padding_m_splits = _pad_tensor_for_fp8(
-                inp_hidden_states, m_splits
-            )
-            padded_inp_hidden_states = block(padded_inp_hidden_states, padding_m_splits)
-            out = _unpad_tensor_for_fp8(padded_inp_hidden_states, m_splits, padding_m_splits)
+            if fp8:
+                padded_inp_hidden_states, padding_m_splits = _pad_tensor_for_fp8(
+                    inp_hidden_states, m_splits
+                )
+                padded_inp_hidden_states = block(padded_inp_hidden_states, padding_m_splits)
+                out = _unpad_tensor_for_fp8(padded_inp_hidden_states, m_splits, padding_m_splits)
+            else:
+                out = block(inp_hidden_states, m_splits)
 
     loss = out.sum()
     loss.backward()