Updated to work with current pytorch

[adishourya]

I have used and tested this with my redweb dataset (rgb iamge and depth map pair) seemed to work fine.

I am open to do more tests on this if you need me to

Best, Aditya