From the error in the colab, you use .data[0]. .data should not be used anymore.
If the goal here is to get a python number. you should replace .data[0] by just .item() (as suggested in the error message).
Many thanks @albanD but after doing this I am facing this error,
RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation: [torch.FloatTensor [32, 16]], which is output 0 of SigmoidBackward, is at version 10; expected version 0 instead. Hint: enable anomaly detection to find the operation that failed to compute its gradient, with torch.autograd.set_detect_anomaly(True).
one of the variables needed for gradient computation has been modified by an inplace operation
This means that you did an inplace operation on a Tensor whose original value was required to be able to compute gradients.
You need to remove this inplace operation to be able to compute gradients.
which is output 0 of SigmoidBackward,
The problematic Tensor is the output of Sigmoid.
So you need to find where you do an inplace operation on the output of the sigmoid and replace that with an out of place op. Or .clone() the output of the sigmoid.