Requires for torch.tensor before casting (#31755)

dc72fd7e · Ella Charlaix · GitHub · 7f91f168 · dc72fd7e
Unverified Commit dc72fd7e authored Jul 03, 2024 by Ella Charlaix Committed by GitHub Jul 03, 2024
Show whitespace changes
Inline Side-by-side

Showing with 2 additions and 2 deletions

src/transformers/utils/generic.py src/transformers/utils/generic.py +2 -2

No files found.
--- a/src/transformers/utils/generic.py
+++ b/src/transformers/utils/generic.py
@@ -762,7 +762,7 @@ def torch_int(x):

    import torch

-    return x.to(torch.int64) if torch.jit.is_tracing() else int(x)
+    return x.to(torch.int64) if torch.jit.is_tracing() and isinstance(x, torch.Tensor) else int(x)


 def torch_float(x):
@@ -774,7 +774,7 @@ def torch_float(x):

    import torch

-    return x.to(torch.float32) if torch.jit.is_tracing() else int(x)
+    return x.to(torch.float32) if torch.jit.is_tracing() and isinstance(x, torch.Tensor) else int(x)


 def filter_out_non_signature_kwargs(extra: Optional[list] = None):