It seems that colab switched to torch 1.6 at the time between my question and your reply.
At first glance I am also no longer able to reproduce the error with torch 1.6 in the colab.
I am still able to reproduce the bug with torch 1.4 with copy pasting the first cell to my local machine (deleting %matplotlib inline and adding as last line plt.show())
I don’t have torch 1.5 installed on my local machine but the code worked yesterday in google colab which was using torch 1.5.
I can’t use

!pip install torch==1.5.0 torchvision==0.6.0

in colab because torch then complains about

The NVIDIA driver on your system is too old (found version 10010). Please update your GPU driver by downloading and installing a new version from the URL: http://www.nvidia.com/Download/index.aspx Alternatively, go to: https://pytorch.org to install a PyTorch version that has been compiled with your version of the CUDA driver.

Are you aware of any fixes in torch 1.6 that might have fixed this error?
EDIT I think I found it myself:
In the torch 1.6 bugfixes there is this line
and I think this fits perfectly for my error

nn.MaxPool3d: fixed incorrect CUDA backward results for non-square output (#36820)


  
      github.com/pytorch/pytorch
  
  
    
  
    
  

  
    
      Fix MaxPool3d CUDA backward incorrect results for non-square output 4
    

    
      pytorch:master ← linziyi96:master
    

    
      
        opened 
        
          
        
        Apr 17, 2020
      
      

      
        
          
          linziyi96
        
      

      
        
          +11
          -1