.backward() on MSELoss fails with IndexError: Dimension out of range on MPS #79784
Description
🐛 Describe the bug
import torch
import torch.nn as nn
loss_fn = torch.nn.MSELoss(reduction="sum")
model = nn.Sequential(
nn.Linear(2, 2)
)
# Works on CPU
x = torch.tensor([1.0, 1.0], dtype=torch.float32)
y = torch.tensor([2.0, 2.0], dtype=torch.float32)
model.zero_grad()
y_predicted = model(x)
loss = loss_fn(y_predicted, y)
loss.backward()
# Doesn't on MPS
device = torch.device("mps")
x_mps = x.to(device=device)
y_mps = y.to(device=device)
model_mps = model.to(device)
model_mps.zero_grad()
y_predicted = model_mps(x_mps)
loss = loss_fn(y_predicted, y_mps)
loss.backward()---------------------------------------------------------------------------
IndexError Traceback (most recent call last)
Input In [48], in <cell line: 25>()
23 y_predicted = model_mps(x_mps)
24 loss = loss_fn(y_predicted, y_mps)
---> 25 loss.backward()
File /opt/homebrew/lib/python3.9/site-packages/torch/_tensor.py:400, in Tensor.backward(self, gradient, retain_graph, create_graph, inputs)
391 if has_torch_function_unary(self):
392 return handle_torch_function(
393 Tensor.backward,
394 (self,),
(...)
398 create_graph=create_graph,
399 inputs=inputs)
--> 400 torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
File /opt/homebrew/lib/python3.9/site-packages/torch/autograd/__init__.py:173, in backward(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)
168 retain_graph = create_graph
170 # The reason we repeat same the comment below is that
171 # some Python versions print out the first line of a multi-line function
172 # calls in the traceback and some print out the last line
--> 173 Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
174 tensors, grad_tensors_, retain_graph, create_graph, inputs,
175 allow_unreachable=True, accumulate_grad=True)
IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
Versions
bash-5.1$ python3 collect_env.py
Collecting environment information...
PyTorch version: 1.13.0.dev20220614
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A
OS: macOS 12.4 (arm64)
GCC version: Could not collect
Clang version: 13.1.6 (clang-1316.0.21.2.5)
CMake version: version 3.23.2
Libc version: N/A
Python version: 3.9.13 (main, May 24 2022, 21:13:51) [Clang 13.1.6 (clang-1316.0.21.2)] (64-bit runtime)
Python platform: macOS-12.4-arm64-arm-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True
Versions of relevant libraries:
[pip3] numpy==1.23.0rc2
[pip3] torch==1.13.0.dev20220614
[pip3] torchaudio==0.14.0.dev20220603
[pip3] torchvision==0.14.0.dev20220614
[conda] Could not collect
cc @ezyang @albanD @zou3519 @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7 @kulinseth