fix

Signed-off-by: Pawel Gadzinski <[email protected]>

Author: pggPL

transformer_engine/debug/features/utils/stats_buffer.py

@@ -88,7 +88,7 @@ def feed(self, tensor, iteration):
         if (
             tensor.numel() == 0
             if hasattr(tensor, "numel")
-            else all(t.numel() == 0 for t in tensor.get_data_tensors())
+            else all((t is None or t.numel() == 0) for t in tensor.get_data_tensors())
         ):
             return
 

transformer_engine/debug/features/utils/stats_computation.py

@@ -127,7 +127,7 @@ def _get(buffers, stat_name):
         lambda buffers: min(_get(buffers, "dynamic_range_bottom")),
     ),
     "underflows_num": (
-        lambda x: (x._data == 0).sum(),
+        lambda x: (x.get_data_tensors()[0] == 0).sum(),
         lambda buffers: sum(_get(buffers, "underflows_num")),
     ),
     "std": (