Support torch convert_to_numpy for all devices (keras-team#20042)

* Support torch convert_to_numpy for all devices * reformat
SamanehSaadat · Jul 25, 2024 · 08e7394 · 08e7394
1 parent c9bd4a8
commit 08e7394
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/keras/src/backend/torch/core.py b/keras/src/backend/torch/core.py
@@ -240,7 +240,7 @@ def transform(x):
  if x.requires_grad:
  x = x.detach()
  # Tensor has to be moved to CPU before converting to numpy.
- if x.is_cuda or x.is_mps:
+ if x.device != torch.device("cpu"):
  x = x.cpu()
  if x.dtype == torch.bfloat16:
  # Attempting to call .numpy() on a bfloat16 torch tensor leads