Try this,
from torch.nn.utils import spectral_norm

def __init__(self, dim=DIM, output_dim=OUTPUT_DIM):
        super(Generator, self).__init__()

        self.dim = dim
        self.ln1 = spectral_norm(nn.Linear(128, 4*4*8*self.dim))


I don’t know how SpectralNorm is defined in your code. Can you check it?