I combined traditional UNET module + my own custom function in my forward model. For some reason, I can’t give data as a batch tensor, but I need to give one by one using forloop.
____________________________________________
optvars = [{'params': custom_var, 'lr':lr_custom}]
optimizer1 = optim.Adam(optvars)
optimizer2 = optim.Adam(UNET.parameters(), lr=lr_UNET)

for a in range(batch_size)
     input_data=batch[a,:,:]
     temp_data=custom_function(input_data, custom_var)
     output_data=UNET(temp_data)
     loss_total=loss_total + loss(output_data)

loss_total.backward()  #  <-----where the error occur
custom_var.retain_grad() # <----- I think I need something like this for  UNET parameters...
optimizer1.step()
optimizer2.step()
____________________________________

I get this error:
RuntimeError: Unable to find a valid cuDNN algorithm to run convolution
It works when batch_size=1, but gives an error when batch_size>=2.
Also, when I don’t use UNET module in the forloop (only using custom_function), it works.
I think it’s something to do with using module more than twice before backward propagation, but I’m not sure what’s the problem.
Can anyone help please?
Thanks!