Also seeing this with no obvious reason why:
[W accumulate_grad.h:185] Warning: grad and param do not obey the gradient layout contract. This is not an error, but may impair performance.
grad.sizes() = [12, 48, 1, 1], strides() = [48, 1, 1, 1]
param.sizes() = [12, 48, 1, 1], strides() = [48, 1, 48, 48] (function operator())

Looks very much to be a regression bug in PyTorch (running 1.9.1 here).