fix bug of ut weight convert

Author: z00573959

tests/modeling_test_utils.py

@@ -9,7 +9,7 @@
 from ml_dtypes import bfloat16
 
 import mindspore as ms
-from mindspore import mint, nn, ops
+from mindspore import nn, ops
 
 logger = logging.getLogger("ModelingsUnitTest")
 
@@ -131,10 +131,6 @@ def get_pt2ms_mappings(m):
                 mappings[f"{name}.running_mean"] = f"{name}.moving_mean", lambda x: x
                 mappings[f"{name}.running_var"] = f"{name}.moving_variance", lambda x: x
                 mappings[f"{name}.num_batches_tracked"] = None, lambda x: x
-        elif isinstance(cell, (mint.nn.BatchNorm1d, mint.nn.BatchNorm2d, mint.nn.BatchNorm3d)):
-            # TODO: for mint.nn, the dtype for each param should expected to be same among torch and mindspore
-            # this is a temporary fix, delete this branch in future.
-            mappings[f"{name}.num_batches_tracked"] = f"{name}.num_batches_tracked", lambda x: x.to(ms.float32)
     return mappings
 
 
@@ -150,6 +146,11 @@ def convert_state_dict(m, state_dict_pt):
     state_dict_ms = {}
     for name_pt, data_pt in state_dict_pt.items():
         name_ms, data_mapping = mappings.get(name_pt, (name_pt, lambda x: x))
+        # for torch back compatibility
+        # for torch <2.0, dtype of num_batches_tracked is int32, for torch>=2.0, dtype of num_batches_tracked is int64,
+        # mindspore.mint is aligned with torch>=2.0
+        if "num_batches_tracked" in name_pt and data_pt.dtype == torch.int32:
+            data_pt = data_pt.to(torch.int64)
         data_ms = ms.Parameter(
             data_mapping(ms.Tensor.from_numpy(data_pt.float().numpy()).to(dtype_mappings[data_pt.dtype])), name=name_ms
         )