From a0715f388d36a69d16433e2fa6bcf1397444f02a Mon Sep 17 00:00:00 2001
From: Diego Bellani <34714625+l0r3m1psum@users.noreply.github.com>
Date: Tue, 10 Sep 2024 09:57:48 +0200
Subject: [PATCH] Added a fix to make the module work inside a virtual
 environment on Windows

On Windows when trying to load the extension, after having compiled it and installed it in a virtual environnement, it failed to load the custom operator. I have traced this issue to a problem related to DLL loading paths.

This change fixes this issue by adding the directory, in the virtual environment, where PyTorch's DLLs are stored, to the DLL search path of Python.
---
 extension_cpp/__init__.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/extension_cpp/__init__.py b/extension_cpp/__init__.py
index 769c697..bbcf859 100644
--- a/extension_cpp/__init__.py
+++ b/extension_cpp/__init__.py
@@ -1,2 +1,6 @@
 import torch
+inside_virtual_environment = sys.prefix != sys.base_prefix
+if inside_virtual_environment and os.name == 'nt':
+    dll_dir = os.path.join(sys.prefix, 'Lib\\site-packages\\torch\\lib')
+    handle = os.add_dll_directory(dll_dir)
 from . import _C, ops