Ok, I can reproduce your results. It turns out that `torch.square`

works in the notebook, but not for the grader. I tried implementing that function three different ways:

```
#### START CODE HERE ####
# penalty = torch.mean(torch.square(gradient_norm - 1))
penalty = torch.mean(torch.pow(gradient_norm - 1,2))
# penalty = torch.mean((gradient_norm - 1)**2)
#### END CODE HERE ####
```

All three of them work in the notebook, but only the `torch.pow`

and `**2`

versions will pass the grader.

So it must be the case that the grader is using a different version of `torch`

that somehow does not support the `square`

function. I checked the current Torch documentation and it doesn’t say anything about “square” being deprecated or the like, so this is something of a mystery. I will try to either file a bug or (failing that) at least report this to the Course Staff in case they think it’s something worth figuring out.