This is what I am trying to run now. Without .cuda() it runs OK.
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.distributions import Uniform, Normal
import torch.utils.data as tdata
import torch.optim as optim
import numpy as np


class DepthToSpace(nn.Module):
    def __init__(self, block_size):
        super().__init__()
        self.block_size = block_size
        self.block_size_sq = block_size * block_size
 
    def forward(self, input):
        output = input.permute(0, 2, 3, 1)
        (batch_size, d_height, d_width, d_depth) = output.size()
        s_depth = int(d_depth / self.block_size_sq)
        s_width = int(d_width * self.block_size)
        s_height = int(d_height * self.block_size)
        t_1 = output.reshape(batch_size, d_height, d_width, self.block_size_sq, s_depth)
        spl = t_1.split(self.block_size, 3)
        stack = [t_t.reshape(batch_size, d_height, s_width, s_depth) for t_t in spl]
        output = torch.stack(stack, 0).transpose(0, 1).permute(0, 2, 1, 3, 4).reshape(batch_size, s_height, s_width,
                                                                                      s_depth)
        output = output.permute(0, 3, 1, 2)
        return output

# Spatial Upsampling with Nearest Neighbors
class Upsample_Conv2d(nn.Module):
  def __init__(self, in_dim, out_dim, kernel_size=(3, 3), stride=1, padding=1):
    super(Upsample_Conv2d, self).__init__()
    self.depth_to_space = DepthToSpace(block_size=2)
    self.conv2d = nn.Conv2d(in_dim, out_dim, kernel_size, stride=stride, 
                            padding=padding)

  def forward(self, x):
    x_ = torch.cat([x, x, x, x], dim=1)
    x_ = self.depth_to_space(x_)
    x_ = self.conv2d(x_)
    return x_


class ReshapeGan(nn.Module):
  def __init__(self, out_shape):
    super(ReshapeGan, self).__init__()
    self.out_shape = out_shape
  
  def forward(self, x):
    b = x.shape[0]
    return x.view(b, *self.out_shape)

class Generator(nn.Module):
  def __init__(self, n_filters=128):
    super(Generator, self).__init__()
    self.n_filters = n_filters
    self.latent = torch.distributions.Normal(torch.tensor(0.), torch.tensor(1.))
    self.network = nn.Sequential(
        nn.Linear(n_filters, 4*4*256),
        ReshapeGan((256,4,4)),
        ResnetBlockUp(in_dim=256, n_filters=n_filters),
        ResnetBlockUp(in_dim=n_filters, n_filters=n_filters),
        ResnetBlockUp(in_dim=n_filters, n_filters=n_filters),
        nn.BatchNorm2d(n_filters),
        nn.ReLU(),
        nn.Conv2d(n_filters, 3, kernel_size=(3, 3), padding=1),
        nn.Tanh()
    )

  def forward(self, z):
    return self.network(z)

  def sample(self, n):
    z = self.latent.sample([n, self.n_filters])
    return self.forward(z)

CUDA_LAUNCH_BLOCKING=1
g = Generator(128)
g = g.cuda() # HERE IT CRASHES
z = g.sample(10)
print(z.shape)