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))))