Replace PyLong_From with THPUtils_packUInt() to avoid overflow errors

pytorch · May 24, 2023 · 944f2df · 944f2df
1 parent 539fac5
commit 944f2df
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/torch/csrc/mps/Module.cpp b/torch/csrc/mps/Module.cpp
@@ -105,7 +105,7 @@ static PyObject* MPSModule_currentAllocatedMemory(
     PyObject* _unused,
     PyObject* noargs) {
   HANDLE_TH_ERRORS
-  return PyLong_FromUnsignedLongLong(
+  return THPUtils_packUInt64(
       at::detail::getMPSHooks().getCurrentAllocatedMemory());
   END_HANDLE_TH_ERRORS
 }
@@ -114,7 +114,7 @@ static PyObject* MPSModule_driverAllocatedMemory(
     PyObject* _unused,
     PyObject* noargs) {
   HANDLE_TH_ERRORS
-  return PyLong_FromUnsignedLongLong(
+  return THPUtils_packUInt64(
       at::detail::getMPSHooks().getDriverAllocatedMemory());
   END_HANDLE_TH_ERRORS
 }
@@ -149,7 +149,7 @@ static PyObject* MPSModule_profilerStopTrace(
 static PyObject* MPSModule_acquireEvent(PyObject* _unused, PyObject* args) {
   HANDLE_TH_ERRORS
   const bool enable_timing = THPUtils_unpackBool(args);
-  return PyLong_FromUnsignedLong(
+  return THPUtils_packUInt32(
       at::detail::getMPSHooks().acquireEvent(enable_timing));
   END_HANDLE_TH_ERRORS
 }