Dear Patrick,
I have created a new conda environement and I’ve been able to install the latest version of Pytorch.
That said, I still had an error related to illegal memory access. I found in another post that it could be related to having variables both on gpu and cpu. So, I checked up on the variable (i.e. weights) created in MyConv2d and it appeared that it was created by default on  cpu. To correct I added the following lines at the top of the forward function :
        if x.is_cuda:
            self.weights = self.weights.cuda()

I also had problem of exploding gradient at the very beginning of training, so I normalized the output of MyConv2d to check if it improves things. For now, this issue seems solved, but I still have a problem with training but this time with memory. The training goes well for a big number of epochs (~100) , but at some point, I get this error :

RuntimeError: CUDA out of memory. Tried to allocate 656.00 MiB (GPU 0; 5.00 GiB total capacity; 345.42 MiB already allocated; 590.35 MiB free; 1.10 GiB reserved in total by PyTorch)

And I don’t understand why it occurs at this point; memory should be free at the end of each epoch, right? Then I have another question : is there a tool (like the Matlab profile viewer) to check the memory consumption for pytorch code?
I post here the last version of my custom convolutional layer, if it helps.
class MyConv2d(nn.Module):
    def __init__(self, n_channels, out_channels, kernel_size, dilation=1, padding=0, stride=1):
        super(MyConv2d, self).__init__()

        self.kernel_size = (kernel_size, kernel_size)
        self.kernel_size_number = kernel_size * kernel_size
        self.out_channels = out_channels
        self.dilation = (dilation, dilation)
        self.padding = (padding, padding)
        self.stride = (stride, stride)
        self.n_channels = n_channels
        self.weights = nn.Parameter(torch.Tensor(self.out_channels, self.n_channels, self.kernel_size_number)).data.uniform_(0, 1)

    def forward(self, x, mask):
        if x.is_cuda:
            self.weights = self.weights.cuda()
        width = self.calculateNewWidth(x)
        height = self.calculateNewHeight(x)
        result = torch.zeros(
            [x.shape[0] * self.out_channels, width, height], dtype=torch.float32, device=device
        )
        result_mask = torch.zeros(
            [x.shape[0] * self.out_channels, width, height], dtype=torch.float32, device=device
        )
        
        
        windows_x = self.calculateWindows(x)
        windows_mask = self.calculateWindows(mask)
        windows_mask[windows_mask < 1] = -1
        windows_mask_centers = windows_mask[:, :, windows_mask.size()[2]//2].view(windows_mask.size()[0], windows_mask.size()[1], 1)
        windows_mask = windows_mask * windows_mask_centers
        windows_mask[windows_mask < 1] = 0 
        windows_x_seg = windows_x * windows_mask
        
        # compute the result of x with mask-aware convolution
        for i_convNumber in range(self.out_channels):
            for channel in range(x.shape[1]):
                xx = torch.matmul(windows_x_seg[channel], self.weights[i_convNumber][channel].view(-1, 1))
                xx = xx.view(-1, width, height)/torch.sum(windows_mask[channel], 1).view(-1, width, height)
                result[i_convNumber * xx.shape[0] : (i_convNumber + 1) * xx.shape[0]] += xx
            result[i_convNumber * xx.shape[0] : (i_convNumber + 1) * xx.shape[0]] /= x.shape[1]

        result = result.view(x.shape[0], self.out_channels, width, height)
 
        # compute the result of mask with mask-aware convolution
        windows_mask_seg = self.calculateWindows(mask) * windows_mask
        for i_convNumber in range(self.out_channels):
            for channel in range(x.shape[1]):
                xx = torch.matmul(windows_mask_seg[channel], self.weights[i_convNumber][channel].view(-1, 1))
                xx = xx.view(-1, width, height)
                result_mask[i_convNumber * xx.shape[0] : (i_convNumber + 1) * xx.shape[0]] += xx

        result_mask = result_mask.view(mask.shape[0], self.out_channels, width, height)
        result_mask = torch.clamp(result_mask, min=0, max=1)
        
        return result, result_mask

    def calculateWindows(self, x):
        windows = F.unfold(
            x, kernel_size=self.kernel_size, padding=self.padding, dilation=self.dilation, stride=self.stride
        )

        windows = windows.transpose(1, 2).contiguous().view(-1, x.shape[1], self.kernel_size_number)
        windows = windows.transpose(0, 1)

        return windows

    def calculateNewWidth(self, x):
        return (
            (x.shape[2] + 2 * self.padding[0] - self.dilation[0] * (self.kernel_size[0] - 1) - 1)
            // self.stride[0]
        ) + 1

    def calculateNewHeight(self, x):
        return (
            (x.shape[3] + 2 * self.padding[1] - self.dilation[1] * (self.kernel_size[1] - 1) - 1)
            // self.stride[1]
        ) + 1

Many thanks in advance