Move the .cuda() call into the nn.Parameter() call:

self.B_conv1 = nn.Parameter(params.items()[1][1].cuda())


If you don’t want to backprop through the dictionary it’s easy:

def linear_decomp(input, coefs, dictionary, bias=None):
   weight = torch.mm(dictionary.detach(), coefs).cuda()
   output = input.mm(weight.t())
   if bias is not None:
     output += bias.unsqueeze(0)
   return output