Hi @ptrblck, yes you are absolutely correct. The issue is visible and random.
I did check the tensors before and after applying the transformation. The tensor before the transformation has no NaNs. Interestingly, I opened the debug console and applied transformation and there was no NaN. But the same transformation produce NaN randomly when running.
Other than that, the same tensor does not give any NaN in another run of the same code. I am unable to comprehend it.
Do you think the issue could be because of reading two datasets of different sizes (though it should not be the case)? Or could you recommend anything else?