I’ve upgraded PyTorch, but sadly the bug persists.
Here is the output of python -m torch.utils.collect_env
PyTorch version: 1.9.0+cu102
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.2 LTS (x86_64)
GCC version: (Ubuntu 9.3.0-17ubuntu1~20.04) 9.3.0
Clang version: 10.0.0-4ubuntu1 
CMake version: Could not collect
Libc version: glibc-2.10

Python version: 3.7 (64-bit runtime)
Python platform: Linux-5.4.0-73-generic-x86_64-with-debian-bullseye-sid
Is CUDA available: True
CUDA runtime version: Could not collect
GPU models and configuration: GPU 0: GeForce RTX 2080 Ti
Nvidia driver version: 460.73.01
cuDNN version: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.5.0
HIP runtime version: N/A
MIOpen runtime version: N/A

Versions of relevant libraries:
[pip] backpack-for-pytorch==1.2.0
[pip] numpy==1.20.2
[pip] pytorch-memlab==0.2.3
[pip] torch==1.9.0
[pip] torchvision==0.10.0
[conda] backpack-for-pytorch      1.2.0                     dev_0    <develop>
[conda] numpy                     1.20.2                   pypi_0    pypi
[conda] pytorch-memlab            0.2.3                    pypi_0    pypi
[conda] torch                     1.9.0                    pypi_0    pypi
[conda] torchvision               0.10.0                   pypi_0    pypi