Ups, sorry:
class ResnetBlockUp(nn.Module):
  def __init__(self, in_dim, kernel_size=(3, 3), n_filters=256):
    super(ResnetBlockUp, self).__init__()
    self.layers = nn.ModuleList(
        [nn.BatchNorm2d(in_dim), 
         nn.ReLU(),
         nn.Conv2d(in_dim, n_filters, kernel_size, padding=1),
         nn.BatchNorm2d(n_filters), 
         nn.ReLU(),
         Upsample_Conv2d(n_filters, n_filters, kernel_size=kernel_size, padding=1)])
    self.shortcut = Upsample_Conv2d(in_dim, n_filters, 
                                    kernel_size=(1, 1), padding=0)

  def forward(self, x):
    x_ = x
    for layer in self.layers:
      x_ = layer(x_)
    x = self.shortcut(x)
    return x + x_

The strange thing is that it crashes with error code:
RuntimeError: cuDNN error: CUDNN_STATUS_NOT_INITIALIZED
right after resetting the colab at the g.sample(10) line. Then at .cuda() line I get the illegal access error after retrying to run without resetting the colab…