[BugFix,Feature] Optional non_blocking in set, to_module and update (#718)

(cherry picked from commit 2dc0285)

Author: vmoens

setup.py

@@ -65,9 +65,7 @@ def write_version_file(version):
 def _get_pytorch_version(is_nightly):
     # if "PYTORCH_VERSION" in os.environ:
     #     return f"torch=={os.environ['PYTORCH_VERSION']}"
-    if is_nightly:
-        return "torch>=2.3.0.dev"
-    return "torch>=2.2.1"
+    return "torch==2.2.2"
 
 
 def _get_packages():