torchvision.ops DeformConv2d creates RuntimeError: CUDA error: an illegal memory access was encountered on CUDA and not on CPU

🐛 Describe the bug

torchvision ops DeformConv2d creates a illegal memory access issue while running on cuda while it works fine on CPU.

setup:-
torch '2.5.1+cu124'
torchvision '0.20.1+cu124'
Python 3.9.11 64 bit

sample code

import torch
import torch.nn as nn
from torchvision.ops import DeformConv2d  

class DeformConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, offset ,kernel_size=3, padding=1):
        super(DeformConvBlock, self).__init__()
        self.offset = offset
        self.deform_conv = DeformConv2d(in_channels, out_channels , kernel_size=kernel_size, padding=padding)

    def forward(self, x):
        return self.deform_conv(x, offset = self.offset)

class DecoderModule(nn.Module):
    def __init__(self):
        super(DecoderModule, self).__init__()
        self.offset1 = torch.rand(1, 2 * 3 * 3, 128, 128).requires_grad_(True)
        self.offset2 = torch.rand(1, 2 * 3 * 3, 128, 128).requires_grad_(True)
        self.offset3 = torch.rand(1, 2 * 3 * 3, 128, 128).requires_grad_(True)

        self.main = nn.Sequential(
            DeformConvBlock(1, 2 ,offset = self.offset1 , kernel_size=3, padding=1),
            nn.BatchNorm2d(2),# Deformable conv
            nn.ReLU(),

            DeformConvBlock(2, 2 ,offset = self.offset2, kernel_size=3, padding=1),
            nn.BatchNorm2d(2),
            nn.ReLU(),

            DeformConvBlock(2, 1, offset = self.offset3 ,  kernel_size=3, padding=1),

        )

    def forward(self, x):
        return self.main(x)

import os
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
os.environ["CUDA_LAUNCH_BLOCKING"] = "1"
r = DecoderModule().to(device) 
r(torch.rand(1,1,128,128).to(device))

Error:-

RuntimeError: CUDA error: an illegal memory access was encountered
Compile with `TORCH_USE_CUDA_DSA` to enable device-side assertions.

Folks please help :)))

Versions

PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Microsoft Windows 11 Home (10.0.26100 64-bit)
GCC version: Could not collect
Clang version: Could not collect
CMake version: Could not collect
Libc version: N/A

Python version: 3.9.11 (tags/v3.9.11:2de452f, Mar 16 2022, 14:33:45) [MSC v.1929 64 bit (AMD64)] (64-bit runtime)
Python platform: Windows-10-10.0.26100-SP0
Is CUDA available: True
CUDA runtime version: 12.6.77
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA GeForce RTX 4070 Laptop GPU
Nvidia driver version: 560.94
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Name: AMD Ryzen 9 7945HX with Radeon Graphics
Manufacturer: AuthenticAMD
Family: 107
Architecture: 9
ProcessorType: 3
DeviceID: CPU0
CurrentClockSpeed: 2501
MaxClockSpeed: 2501
L2CacheSize: 16384
L2CacheSpeed: None
Revision: 24834

Versions of relevant libraries:
[pip3] numpy==2.0.2
[pip3] torch==2.5.1+cu124
[pip3] torchaudio==2.5.1+cu124
[pip3] torchvision==0.20.1+cu124
[conda] Could not collect

Start by running the provided DeformConv2d sample with CUDA_LAUNCH_BLOCKING=1 and compare it with the CPU path. Then inspect the torchvision.ops.DeformConv2d entry point and its CUDA implementation to isolate the failing input or operation. Done means a confirmed cause and a fix or a precise report showing the supported behavior.