Gradient of using the output of the network twice

Below is a toy code.
(The actual model and USERFUNC are more complex.)

def USERFUNC(network_output):
  comp = network_output + 1j*network_output
  comp = comp * 123
  func_out = torch.abs(comp)
  return func_out

model = nn.Linear(1,1)

criterion = nn.MSELoss()

x = torch.randn(1,1)

out = model(x)

out_c = out.clone()

lossA = criterion(out, x)

out_prime = USERFUNC(out_c)

lossB = criterion(out_prime, x)

a = 0.5
loss = lossA + a * lossB


The toy code seems to work normally, but the actual learning result seems strange.
I think the cause is related to gradient.

If only loss A is applied, it will operate normally, but if loss B is applied, the result will be wrong.
I thought USERFUNC was wrong, but this should not be a problem because I checked that the input of USERFUNC was operated with a = torch.randn (3,256,256, device='cuda', requirements_grad=True).

What is the problem? I’ve been wandering this problem for months now…:frowning:
Please let me know if you need more explanation.

Any response would be appreciated.
Thank you.

What comp = network_output + 1j*network_output this line does?

Thanks for your reply!

The line is to make it complex number.
(It may seem strange, but I need it…)

here 1 is constant and j is a variable?

1 is constant, j is a imaginary i.

replace line

func_out = comp.abs()

with this

func_out = torch.abs(comp)

Oh, that’s my typo. I’m already using func_out = torch.abs(comp) as you suggest.
I’m sorry to confuse you.:sob:

What is the final code and error you are receiving?

Sorry for replying late.

The error does not occur. It’s just that the learning results are weird.
(There is no final code, but the toy code is the overall framework.)

The meaning of “weird” means that it seems to return the loss to be average as a result of learning.
For example, if the probability model returns 0.5