@albanD sure I’m just working on some pixelCNN tutorial myself so I can share my full code here.
class MaskedConv3d(nn.Conv3d):
    def __init__(self, mask_type, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.register_buffer("mask", self.weight.clone())
        out_channel, in_channel, h, w, c = self.weight.size()
        self.mask.fill_(1)
        self.mask[:, :, h // 2: , w // 2 + 1:] = 0
        self.mask[:, :, h // 2 + 1:] = 0
        self.mask[:, :, h // 2, w // 2, c // 2 + (mask_type == 'B')] = 0
        
    def forward(self, x):
        with torch.no_grad():
            self.weight *= self.mask
        return super().forward(x)

net = nn.Sequential(
    MaskedConv3d('A', 1, fm, (7, 7, 3), 1, (3, 3, 1), bias=False), nn.BatchNorm3d(fm), nn.ReLU(inplace=True),
    MaskedConv3d('B', fm, fm, (7, 7, 3), 1, (3, 3, 1), bias=False), nn.BatchNorm3d(fm), nn.ReLU(inplace=True),
    MaskedConv3d('B', fm, fm, (7, 7, 3), 1, (3, 3, 1), bias=False), nn.BatchNorm3d(fm), nn.ReLU(inplace=True),
    MaskedConv3d('B', fm, fm, (7, 7, 3), 1, (3, 3, 1), bias=False), nn.BatchNorm3d(fm), nn.ReLU(inplace=True),
    MaskedConv3d('B', fm, fm, (7, 7, 3), 1, (3, 3, 1), bias=False), nn.BatchNorm3d(fm), nn.ReLU(inplace=True),
    MaskedConv3d('B', fm, fm, (7, 7, 3), 1, (3, 3, 1), bias=False), nn.BatchNorm3d(fm), nn.ReLU(True),
    MaskedConv3d('B', fm, fm, (7, 7, 3), 1, (3, 3, 1), bias=False), nn.BatchNorm3d(fm), nn.ReLU(True),
    MaskedConv3d('B', fm, fm, (7, 7, 3), 1, (3, 3,1 ), bias=False), nn.BatchNorm3d(fm), nn.ReLU(True),
    nn.Conv3d(fm, 256 , 1),
)

how do I know if the first op is in-place or not?
is the self.mask.fill_ causing issues?
why isn’t in-place operation allowed here? because you cannot do that on multi-gpu because you cannot change values in place for multi-gpus?