I’m not familiar with your model, but this line of code:
p = torch.cat([w, k, y, z], 0)

would concatenate the different activations in the batch dimension, which is most likely wrong.
If you would like to concatenate the activations instead in the channel dimension, you would have to use dim=1.
Currently this line is changing the batch size, which is probably creating the error.