fix: update pytorch nightly to fix memory leak (#255)

ELS-RD · Jan 20, 2023 · 306aa3a · 306aa3a
1 parent 46e8a58
commit 306aa3a
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/Dockerfile b/Dockerfile
@@ -29,7 +29,7 @@ RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.8 1 &&
 
 RUN python3.9 -m ensurepip --default-pip --upgrade
 
-RUN pip install --pre torch==2.0.0.dev20230104+cu117 --extra-index-url https://download.pytorch.org/whl/nightly/cu117
+RUN pip install --pre torch==2.0.0.dev20230119+cu117 --extra-index-url https://download.pytorch.org/whl/nightly/cu117
 
 
 WORKDIR /syncback

diff --git a/requirements.txt b/requirements.txt
@@ -1,5 +1,5 @@
 triton==2.0.0.dev20221202
-torch==2.0.0.dev20230104+cu117
+torch==2.0.0.dev20230119+cu117
 pytest
 tabulate
 termcolor

diff --git a/src/kernl/optimizer/cuda_graph.py b/src/kernl/optimizer/cuda_graph.py
@@ -17,8 +17,8 @@
 
 import torch
 import triton
+from torch._dynamo import utils as dynamo_utils
 from torch._inductor.compile_fx import cudagraphify_impl
-from torch._inductor.utils import dynamo_utils
 from torch._subclasses import FakeTensor