ValueError: invalid type: 'torch.mps.FloatTensor'

# Error Message With `FocalLoss` on `mps` device `target = target.type(output.type())` fails due to `ValueError: invalid type: 'torch.mps.FloatTensor'` ``` | .....venv/lib/python3.12/site-packages/segmentation_models_pytorch/losses/_functional.py:69 in │ │ focal_loss_with_logits │ │ │ │ 66 │ References: │ │ 67 │ │ https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/loss/losses.py │ │ 68 │ """ │ │ ❱ 69 │ target = target.type(output.type()) │ │ 70 │ │ │ 71 │ logpt = F.binary_cross_entropy_with_logits(output, target, reduction="none") │ │ 72 │ pt = torch.exp(-logpt) │ │ │ │ ╭─────────────────────────────────────────── locals ───────────────────────────────────────────╮ │ │ │ alpha = None │ │ │ │ eps = 1e-06 │ │ │ │ gamma = 2.0 │ │ │ │ normalized = False │ │ │ │ output = tensor([[[-0.0172, -0.1538, -0.2904, ..., 0.4121, 0.4157, 0.4192], │ │ │ │ │ │ [-0.0268, -0.1523, -0.2779, ..., 0.4078, 0.4074, 0.4070], │ │ │ │ │ │ [-0.0364, -0.1509, -0.2653, ..., 0.4034, 0.3991, 0.3947], │ │ │ │ │ │ ..., │ │ │ │ │ │ [ 0.2922, 0.3418, 0.3915, ..., 0.2892, 0.2776, 0.2660], │ │ │ │ │ │ [ 0.3021, 0.3555, 0.4090, ..., 0.2808, 0.2703, 0.2598], │ │ │ │ │ │ [ 0.3120, 0.3692, 0.4265, ..., 0.2724, 0.2630, 0.2535]]], │ │ │ │ │ device='mps:0') │ │ │ │ reduced_threshold = None │ │ │ │ reduction = 'mean' │ │ │ │ target = tensor([[[0, 0, 0, ..., 0, 0, 0], │ │ │ │ │ │ [0, 0, 0, ..., 0, 0, 0], │ │ │ │ │ │ [0, 0, 0, ..., 0, 0, 0], │ │ │ │ │ │ ..., │ │ │ │ │ │ [0, 0, 0, ..., 0, 0, 0], │ │ │ │ │ │ [0, 0, 0, ..., 0, 0, 0], │ │ │ │ │ │ [0, 0, 0, ..., 0, 0, 0]]], device='mps:0') │ │ │ ╰──────────────────────────────────────────────────────────────────────────────────────────────╯ │ ╰─────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────────╯ ValueError: invalid type: 'torch.mps.FloatTensor' ``` # Most Simple reproduce ``` Python 3.12.7 (main, Oct 16 2024, 07:12:08) [Clang 18.1.8 ] on darwin Type "help", "copyright", "credits" or "license" for more information. >>> import torch >>> a = torch.rand(4,4, device="mps") >>> b = torch.rand(5,5,device="cpu") >>> b.type(a.type()) # <--- this is the error Traceback (most recent call last): File "<stdin>", line 1, in <module> ValueError: invalid type: 'torch.mps.FloatTensor' >>> b.to(a.type()) # <--- this is the error Traceback (most recent call last): File "<stdin>", line 1, in <module> RuntimeError: Invalid device string: 'torch.mps.FloatTensor' >>> a.type() 'torch.mps.FloatTensor' >>> a.device device(type='mps', index=0) >>> b.to(a.device) # <--- this is the fix tensor([[0.1040, 0.8790, 0.8257, 0.5309, 0.5873], [0.8089, 0.5193, 0.1906, 0.3800, 0.4528], [0.5679, 0.0395, 0.8139, 0.2062, 0.9756], [0.9427, 0.9405, 0.2378, 0.8071, 0.1869], [0.6196, 0.5224, 0.3225, 0.4819, 0.0783]], device='mps:0') >>> ``` # Suggested Fix ``` --- a/segmentation_models_pytorch/losses/_functional.py +++ b/segmentation_models_pytorch/losses/_functional.py @@ -66,7 +66,7 @@ def focal_loss_with_logits( References: https://github.com/open-mmlab/mmdetection/blob/master/mmdet/core/loss/losses.py """ - target = target.type(output.type()) + target = target.to(dtype=output.dtype, device=output.device) logpt = F.binary_cross_entropy_with_logits(output, target, reduction="none") pt = torch.exp(-logpt) ```