"buffer is not large enough" error

When running on device = 'mps' on my M1 macbook, I get the following error:

/AppleInternal/Library/BuildRoots/a0876c02-1788-11ed-b9c4-96898e02b808/Library/Caches/com.apple.xbs/Sources/MetalPerformanceShaders/MPSCore/Types/MPSNDArray.mm:782: failed assertion `[MPSNDArray, initWithBuffer:descriptor:] Error: buffer is not large enough. Must be 8000 bytes

This seems to occur when running the following code (I had to do this workaround of creating a one-hot encoding with numpy and then converting it back to a torch.Tensor because the built-in pytorch one-hot converter seems not to be supported for device = 'mps' yet):

gen_labels = np.random.randint(0, n_classes, current_batchSize)
cls_one_hot = np.zeros((gen_labels.size, gen_labels.max() + 1))
cls_one_hot[np.arange(gen_labels.size), gen_labels] = 1
gen_labels = torch.Tensor(gen_labels).to(device)
cls_one_hot = torch.Tensor(cls_one_hot).to(device)
cls_one_hot = cls_one_hot.type(torch.LongTensor)

I am on Pytorch version 1.13.0.dev20220713.
Thankful for any help!