My guess is, that this line of code is raising the issue:
# Setup the loss function
criterion = torch.nn.functional.mse_loss()

Since you are using the functional API, you would have to pass the values directly as:
loss = torch.nn.functional.mse_loss(output, target)

If you want to create the criterion as a module, use this:
criterion = nn.MSELoss()

Let me know, if that helps.