Currently you are pushing delta to the default device manually (cuda:0 if you haven’t used a context manager to specify the device).
I assume model_A is wrapped in nn.DataParallel, which might trigger this error, since input will be on the specified devices.
If you need to create a new tensor inside your forward, you should push it to the current device:
delta = torch.randn(1, device=input.device)