To avoid using nn.grad.conv2d_input, nn.grad.conv2d_weight as their performance is problematic in the backward pass, I tried using nn.Conv2d with few modifications. My use case is using only subset of the weight tensor’s input channels and masking before applying convolution. For example,
w_true = nn.Parameter(torch.randn(32, 16, 3, 3))
w_modfd = torch.rand_like(w_true) * w_true  #masking
w_modfd = w_modfd[:, :8] #using only subset of input channels

dummy_conv = nn.Conv2d(1,1,1)

while True:
  input = torch.randn(128, 8, 28, 28)
  dummy_conv.weight.data = w_modfd

  out = dummy_conv(x)

As the nn.Conv2d modules performance is as expected, this runs smoothly and accuracy is not also affected as I see. But are there any obvious pitfalls in this method with regard to memory usage, correctness and etc.? One thing I noticed is w_true.grad is not calculated and if we require the gradient of w_true, we should access it through dummy_conv.weight.grad.