test cudnn overwrite

pytorch · Jan 26, 2024 · 41ae59b · HolyWu · Jan 26, 2024 · gs-olive
1 parent cd30eae
commit 41ae59b
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/packaging/smoke_test.py b/packaging/smoke_test.py
@@ -1,13 +1,17 @@
+import glob
 import os
 import shutil
 import site
 
 python_site_packages = site.getsitepackages()[-1]
 extra_cudnn = os.path.join(python_site_packages, "nvidia", "cudnn")
-print(f"Deleting directory: {extra_cudnn}")
-shutil.rmtree(extra_cudnn)
+extra_torch = os.path.join(python_site_packages, "torch", "lib")
+
+for dll in glob.glob(os.path.join(extra_cudnn, "**", "*.dll"), recursive=True):
+    shutil.copy(extra_cudnn, extra_torch)
 
 
 # import tensorrt
 import torch
+
 # import torch_tensorrt