@kouyoumin thank you so much for your post - I am trying to replicate with the most recent version of PyTorch and it looks like the spots where you edited the code are different now - any pointers as to where these changes could be made? (really excited to try this on our new NVSwitch’d A100s to see if the unified memory can truly be “one giant GPU” as NVIDIA claims with a contiguous address space…)