dhyey:

what loss function you are using


Can you please elaborate on “the loop in which you are changing probablities to prediction is causing the problem”? Thank you!
I am using the following loss function:
loss_fn = torch.nn.BCELoss()

# define learning rate and optimizer

learning_rate = 1e-4

optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)