fixed epoch looping and kernel on cpu errors

Author: danielryckman

traincheck/config/config.py

@@ -21,6 +21,7 @@
     "torch.fx",
     "torch._dynamo",
     "torch._sources",  # FIXME: cannot handle this module, instrumenting it will lead to exceptions: TypeError: module, class, method, function, traceback, frame, or code object was expected, got builtin_function_or_method
+    "torchtitan.trainer.Trainer.train_step",
     # "torch.autocast",
     # "torch.amp",
     # "torch.matmul",

traincheck/instrumentor/tracer.py

@@ -634,6 +634,14 @@ def instrument(self) -> int:
 
         global IS_INSTRUMENTING
         IS_INSTRUMENTING = True
+
+        try:
+            import torch.utils._device as _tc_torch_device
+
+            _tc_torch_device._device_constructors()
+        except Exception:
+            pass
+
         visited_file_paths: set[str] = set()
 
         first_pass_instrumented_count = 0

traincheck/trace/trace_pandas.py

@@ -558,7 +558,9 @@ def get_column_dtype(self, column_name: str) -> type:
             return self.column_dtypes_cached[column_name]
 
         filtered_values = self.events[column_name].dropna()
-        filtered_values = filtered_values[filtered_values != MD_NONE()]
+        #filtered_values = filtered_values[filtered_values != MD_NONE()]
+        if filtered_values.dtype == object:
+            filtered_values = filtered_values[filtered_values != MD_NONE()]
 
         if filtered_values.empty:
             self.column_dtypes_cached[column_name] = MD_NONE