convert non-cpu const tensor storage to cpu before serialization

manuelcandales · manuelcandales · commit c2eeb962b50c · 2024-12-06T10:01:45.000-05:00
diff --git a/exir/emit/_emitter.py b/exir/emit/_emitter.py
@@ -368,6 +368,8 @@ def _tensor_spec_to_evalue(self, spec: TensorSpec) -> EValue:
         if spec.const:
             # Tensor with a blob we need to serialize. May not actually be constant at runtime
             # if it's a weight with an associated gradient
+            if spec.storage.device.type != "cpu":
+                spec.storage = spec.storage.cpu()
             spec_array_type = (
                 ctypes.c_char * typing.cast(torch.UntypedStorage, spec.storage).nbytes()
             )

Original file line number	Diff line number	Diff line change
`@@ -368,6 +368,8 @@ def _tensor_spec_to_evalue(self, spec: TensorSpec) -> EValue:`
`368`	`368`	`if spec.const:`
`369`	`369`	`# Tensor with a blob we need to serialize. May not actually be constant at runtime`
`370`	`370`	`# if it's a weight with an associated gradient`
	`371`	`+ if spec.storage.device.type != "cpu":`
	`372`	`+ spec.storage = spec.storage.cpu()`
`371`	`373`	`spec_array_type = (`
`372`	`374`	`ctypes.c_char * typing.cast(torch.UntypedStorage, spec.storage).nbytes()`
`373`	`375`	`)`