I wrote a simple example to show the issue below. As soon as the nn.Linear layer receives an input of size greater than 127 say, 128, my kernel restarts. But for smaller values it works fine.
My pytorch version is 1.7.1 and I am working on the Mac M1.
class Network(nn.Module): def __init__(self): super().__init__() # This makes the kernel restart self.fc1 = nn.Linear(128,1) #anything less than this works fine def forward(self, x): x = self.fc1(x) return x model = Network() print(model.forward(torch.rand((1,128))))