Solved by Caruso in post #2 
              
              
                Hi, 
try this: 
A = torch.tensor(A).cuda().requires_grad_(True)  # or torch.tensor(A, requires_grad_=True, device=device) with device="cuda:0" as example
B = torch.tensor(B).cuda()

data_optimizer = torch.optim.Adam([A], lr=1e-3)

C = torch.cat((A, B), 1)

This way we say that only for A the gradien…