Kernel dies when I use a tensor of size greater than 127 in a linear layer (on Mac M1)

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

Could you execute the script in a terminal and check, if you get a valid error message (I assume you are using a Jupyter Notebook at the moment).
If you get an error message, could you post it here, please?

I get the following error:

OMP: Error #15: Initializing libiomp5.dylib, but found libomp.dylib already initialized.

OMP: Hint This means that multiple copies of the OpenMP runtime have been linked into the program. That is dangerous, since it can degrade performance or cause incorrect results. The best thing to do is to ensure that only a single OpenMP runtime is linked into the process, e.g. by avoiding static linking of the OpenMP runtime in any library. As an unsafe, unsupported, undocumented workaround you can set the environment variable KMP_DUPLICATE_LIB_OK=TRUE to allow the program to continue to execute, but that may cause crashes or silently produce incorrect results. For more information, please see Intel® Product Support.

I was able to solve this by installing:

conda install nomkl

1 Like