Also, are you sure you need a custom autograd function for LinearDecomp? Can you just write it as a normal Python function leveraging autograd?
i.e.
def linear_decomp(input, coefs, dictionary, bias=None):
   weight = torch.mm(dictionary, coefs).cuda()
   output = input.mm(weight.t())
   if bias is not None:
     output += bias.unsqueeze(0)
   return output