Fix slicing for 0-dim param (#36580)

SunMarc · fxmarty-amd · web-flow · commit bc3d5781e72d · 2025-03-13T12:16:13.000+01:00
* fix

* switch to ellipsis instead

* Add co-author
Co-authored-by: fxmarty-amd &lt;fxmarty-amd@users.noreply.github.com&gt;

* Add co-author second try
Co-authored-by: fxmarty-amd &lt;felmarty@amd.com&gt;
diff --git a/src/transformers/integrations/tensor_parallel.py b/src/transformers/integrations/tensor_parallel.py
@@ -531,7 +531,7 @@ def shard_and_distribute_module(
             param, empty_param, param_type, param_casting_dtype, is_contiguous, rank, device_mesh
         )
     else:
-        param = param[:]
+        param = param[...]
         if is_contiguous:
             param = param.contiguous()
 
diff --git a/src/transformers/modeling_utils.py b/src/transformers/modeling_utils.py
@@ -795,7 +795,7 @@ def _load_state_dict_into_meta_model(
                 device_mesh,
             )
         else:
-            param = param[:]
+            param = param[...]
             if casting_dtype is not None:
                 param = param.to(casting_dtype)
             if to_contiguous:

Original file line number	Diff line number	Diff line change
`@@ -531,7 +531,7 @@ def shard_and_distribute_module(`
`531`	`531`	`param, empty_param, param_type, param_casting_dtype, is_contiguous, rank, device_mesh`
`532`	`532`	`)`
`533`	`533`	`else:`
`534`		`- param = param[:]`
	`534`	`+ param = param[...]`
`535`	`535`	`if is_contiguous:`
`536`	`536`	`param = param.contiguous()`
`537`	`537`
Original file line number	Diff line number	Diff line change
`@@ -795,7 +795,7 @@ def _load_state_dict_into_meta_model(`
`795`	`795`	`device_mesh,`
`796`	`796`	`)`
`797`	`797`	`else:`
`798`		`- param = param[:]`
	`798`	`+ param = param[...]`
`799`	`799`	`if casting_dtype is not None:`
`800`	`800`	`param = param.to(casting_dtype)`
`801`	`801`	`if to_contiguous:`