My code is a little messy. I attached here…
class DFR(torch.nn.Module):
def init(self, in_features, out_features, bias=False, name =‘DFR’ ):
super(DFR, self).init()
      self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
      if bias:
           self.bias = nn.Parameter(torch.Tensor(out_features))
      else:
           self.register_parameter('bias', None)

      self.weight.data = torch.zeros(in_features, out_features, device=torch.device('cuda'))
      self.weight.require_grad = True

def forward(self, input):
      # the recevied the input includes two parts: (1) the input (2) the delay_res value
      r = torch.chunk(input, 2)
      delay_res_local = r[1]
      input = r[0]

      input = F.relu(input)
      num = torch.numel(input)  # the element number of the input tensor
      input = input.view(num)

      delay_res_local = delay_res_local.view(num)
      delay_res_local = torch.roll(delay_res_local,1,-1)
      delay_res_local = 0.1 * delay_res_local + input


      delay_res_local = delay_res_local
      h1 = delay_res_local.view(200, 512, 4, 4)
      output = F.relu (h1)

      return output, delay_res_local

in the training parts, the error occurs at the moment of trying accessing “param.grad.data”. Then the error is
param.grad.data = wage_quantizer.QG(param.grad.data,args.wl_grad,grad_scale)
AttributeError: ‘NoneType’ object has no attribute ‘data’
Thank you very much