I have a complex model and there is a part in the model which has this class, the output of this class is (2,1280,208,208) and this is sent to a conv2d which produces (2, 10,208,208) and I am trying to do pixel-wise segmentation. When i do loss.backward() I get segmentation fault. I can see the segmentation fault happens in the cat function because when i return just layer1 or any single layer in CrossAttention1 class I do not get segmentation fault.
class CrossAttention1(nn.Module):
    def __init__(self, dim):
        super(CrossAttention1, self).__init__()
        self.up5 = nn.Upsample(scale_factor=8, mode='bilinear')
        self.up4 = nn.Upsample(scale_factor=8, mode='bilinear')
        self.up3 = nn.Upsample(scale_factor=4, mode='bilinear')
        self.up2 =nn.Upsample(scale_factor=2, mode='bilinear')

    def forward(self, layer1, layer2, layer3, layer4, layer5):
        layer2 = self.up2(layer2)
        layer3 = self.up3(layer3)
        layer4 = self.up4(layer4)
        layer5 = self.up5(layer5)
        x = torch.cat([layer1, layer2, layer3, layer4, layer5], 1)
        return x