alexzhou907/butterflyflow

butterflyflow's Issues

CUDA error: invalid configuration argument

Hi,

I am trying to add one more custom layer (custom autograd ) to the module with CUDA code.
when I try to train the model it shows this error

  logdet_accum = logdet_accum + logdet
                     ~~~~~~~~~~~~~^~~~~~~~
  RuntimeError: CUDA error: invalid configuration argument

How can I resolve this issue?
is it possible to run this code without dataparallel?

See attached autograd for the custom layer

    @staticmethod
    def forward(ctx, x, W, output_x) : # inverse pass for the Inv-Flow
        # z = F.conv2d(x, W)
        # torch.cuda.empty_cache()
        # torch.cuda.synchronize()
        ctx.save_for_backward(x, W)

        # output_x = 0.0 * output_x
        # batch, channels, height, width = x.size()
        # print('inv_conv_, fwd: ', x.device, W.device)
        z = inv_conv_with_bp.inverse(x, W, output_x) # y = output
        # ctx.save_for_backward(x, W, z[0])
      
        return z[0]
    
    @staticmethod
    def backward(ctx, output_grad):
        input, kernel = ctx.saved_tensors
        # input = input.resize(1, input.shape[1], input.shape[2])
        # kernel = kernel.resize(kernel.shape[2], kernel.shape[3])
        b, c, n,m = input.shape
        c_out,c_in, k_h, k_w = kernel.shape
        # M_dk = torch.zeros((b, c, k_h,k_w, n,m), dtype=output_grad.dtype()).to(output_grad.device)
        M_dk = torch.zeros((b, c, k_h,k_w, n,m)).to(output_grad.device) # dtype=torch.float64)
        M_dy = torch.zeros_like(input).to(input.device) # torch.float64
        
        output_dy = torch.zeros_like(input).to(input.device)
        
        # input_grad = inv_conv.dy(output_grad, kernel, M_dy, output_dy)
        input_grad = inv_conv_with_bp.dy(output_grad, kernel, M_dy, output_dy)


        output_dk = torch.zeros_like(kernel).to(kernel.device)
        # grad_k = inv_conv.dw(input, kernel, output_grad, M_dk, output_dk)
        grad_k = inv_conv_with_bp.dw(input, kernel, output_grad, M_dk, output_dk)

        return input_grad[0] , grad_k[0]
    def clip_gradients(self):
        for param in self.parameters():
            if param.grad is not None:
                param.grad.data.clamp_(-self.grad_clip_value, self.grad_clip_value)

  pythorch and cuda version 
  pytorch                   2.1.1           py3.11_cuda11.8_cudnn8.7.0_0    pytorch
  pytorch-cuda              11.8                 h7e8668a_5    pytorch
  pytorch-mutex             1.0                        cuda    pytorch
  torchaudio                2.1.1               py311_cu118    pytorch
  torchtriton               2.1.0                     py311    pytorch
  torchvision               0.16.1              py311_cu118    pytorch