@@ -606,0 +615,4 @@
+                for i in range(num_storages):
+                    args = pickle_module.load(f, **pickle_load_args)
+                    key, location, storage_type = args
+                    dtype = storage_type._dtype

@@ -606,0 +632,4 @@
+                    element_size = _utils._element_size(root.dtype)
+                    offset_bytes = offset * element_size
+                    deserialized_objects[target_cdata] = storage.TypedStorage(
+                        wrap_storage=root._untyped_storage[offset_bytes:offset_bytes + numel * element_size],

@@ -606,0 +621,4 @@
+                    data = np.fromfile(f, dtype=np.uint8, count=nbytes, offset=0)
+                    obj = _UntypedStorage.from_buffer(data)
+
+                    deserialized_objects[key] = storage.TypedStorage(

@@ -606,0 +624,4 @@
+                    deserialized_objects[key] = storage.TypedStorage(
+                        wrap_storage=obj,
+                        dtype=dtype,
+                        _internal=True)

@@ -606,0 +634,4 @@
+                    deserialized_objects[target_cdata] = storage.TypedStorage(
+                        wrap_storage=root._untyped_storage[offset_bytes:offset_bytes + numel * element_size],
+                        dtype=root.dtype,
+                        _internal=True)

@@ -606,0 +649,4 @@
+                    stride = struct.unpack(f'<{ndim}q', f.read(8 * ndim))
+                    storage_offset, = struct.unpack('<q', f.read(8))
+                    tmp_tensor = tensor([], dtype=storage.dtype).set_(
+                        storage._untyped_storage, storage_offset, numel, stride)

@@ -621,3 +650,3 @@
 
     def __setstate__(self, state):
-        Tensor.__init__(self, state["input_data"], dtype=state["dtype"], inner=True)
+        if not isinstance(state, tuple):

@@ -201,3 +188,1 @@
-        source = cast_to_ms_tensor(source)
-        self.set_data(source, True)
-        return self
+    #

@@ -202,2 +188,2 @@
-        self.set_data(source, True)
-        return self
+    #
+    # def set_(self, source=None, storage_offset=0, size=None, stride=None):

@@ -593,3 +596,4 @@
+                    pass
             return super().find_class(mod_name, name)
 
     def legacy_load(f):

@@ -598,2 +602,4 @@
         def persistent_load(saved_id):
             if isinstance(saved_id, tuple):
+                if all(saved_id[1:]):
+                    _check_container_source(*saved_id)

@@ -0,0 +4,4 @@
+    if type is Tensor:
+        return func(*args)
+
+    ret = func(*args).as_subclass(type)

@@ -24,10 +24,14 @@ from typing import Any, BinaryIO, Union, IO, Optional, Type, Dict, Tuple
 from typing_extensions import TypeAlias
 from ml_dtypes import bfloat16
 import numpy as np
+from mindtorch.torch import storage

@@ -448,3 +440,1 @@
-            if storage_dataptr != 0:
-                if storage_dataptr in storage_dtypes:
-                    if storage_dtype != storage_dtypes[storage_dataptr]:
+            if storage.data_ptr() != 0:

@@ -550,3 +546,3 @@
         name = f'archive/data/{key}'
         storage = serialized_storages[key]
-        storage_data = storage.get_bytes()
+        storage_data = storage.inner_data

@@ -590,3 +590,1 @@
-                return eval(name)
-            if mod_name == 'torch':
-                return str(name)
+            if type(mod_name) is str and mod_name.startswith('torch.'):

@@ -791,3 +807,1 @@
-                return eval(name)
-            if mod_name == 'torch':
-                return str(name)
+            if type(mod_name) is str and mod_name.startswith('torch.'):

@@ -54,0 +56,4 @@
+    unsupported_attr(stride)
+    from mindtorch.torch.tensor import tensor # pylint: disable=R0401, C0415
+    t = tensor([], dtype=storage.dtype, device=storage._untyped().device)
+    return t.set_(storage._untyped(), storage_offset, size)

@@ -54,0 +67,4 @@
+
+def _rebuild_device_tensor_from_numpy(data, dtype, device, requires_grad):
+    from mindtorch.torch.functional import from_numpy # pylint: disable=R0401, C0415
+    tensor = from_numpy(data).to(dtype=dtype, device=device)

@@ -45,0 +45,4 @@
+    def __new__(cls, data=None, requires_grad=True, name=None, layerwise_parallel=False, # pylint: disable = W0613
+                parallel_optimizer=True): # pylint: disable = W0613
+        if data is None:
+            data = np.array(1)

@@ -202,3 +194,1 @@
-        self.set_data(source, True)
-        return self
-
+    def __setstate__(self, state):

@@ -581,1 +568,3 @@
-
+                torch_pop()
+                return result
+        torch_disable()

@@ -597,0 +587,4 @@
+        try:
+            current_source = ''.join(get_source_lines_and_file(container_type)[0])
+        except Exception:  # saving the source is optional, so we can ignore any errors
+            warnings.warn("Couldn't retrieve source code for container of "

@@ -14,3 +5,2 @@
-    set_mode_by_env_config()
-    test_import()
+#

@@ -185,0 +358,4 @@
+    pytorch.save(state_dict, 'mindtorch_module_oldfile.pt', _use_new_zipfile_serialization=False)
+    os.remove('mindtorch_module_oldfile.pt')
+
+@SKIP_ENV_CPU(reason="This function need torch version >= 2.1.0")

@@ -219,3 +219,3 @@
                  shuffle: Optional[bool] = None, sampler: Union[Sampler, Iterable, None] = None,
                  batch_sampler: Union[Sampler[Sequence], Iterable[Sequence], None] = None,
-                 num_workers: int = 1, collate_fn: Optional[_collate_fn_t] = None,
+                 num_workers: int = 0, collate_fn: Optional[_collate_fn_t] = None,