Awesome work!
Sorry, was too slow to find the issue.
Generally I wouldn’t recommend to use resize_ or resize_as_ to initialize tensors, since new memory is uninitialized and might yield strange results. In my opinion if you just want to initialize some tensors you should use the factory methods (e.g. torch.zeros etc.).
From the docs:

If the number of elements is larger than the current storage size, then the underlying storage is resized to fit the new number of elements. If the number of elements is smaller, the underlying storage is not changed. Existing elements are preserved but any new memory is uninitialized.