This is another case where for some unknown reason, they use a different version of torch in the assignment notebook than in the grader. Apparently the grader uses an older version that does not support torch.square as a function. Why they would do that is a mystery.
There’s another such “versionitis” issue described on this thread having to do with the behavior of torch.cat.