We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
2 parents b9c3f4e + 9e60cdb commit edebe4dCopy full SHA for edebe4d
modules/torch_utils.py
@@ -20,7 +20,6 @@ def get_param(model) -> torch.nn.Parameter:
20
21
def float64(t: torch.Tensor):
22
"""return torch.float64 if device is not mps or xpu, else return torch.float32"""
23
- match t.device.type:
24
- case 'mps', 'xpu':
25
- return torch.float32
+ if t.device.type in ['mps', 'xpu']:
+ return torch.float32
26
return torch.float64
0 commit comments