From b3213b221d33626b938821a4ded5cdb72f49035c Mon Sep 17 00:00:00 2001
From: Lev Kokotov <hi@levthe.dev>
Date: Mon, 7 Aug 2023 16:53:14 -0700
Subject: [PATCH] Fix clear_gpu_cache

---
 pgml-extension/requirements.txt                 | 1 +
 pgml-extension/src/bindings/transformers/mod.rs | 2 ++
 2 files changed, 3 insertions(+)

diff --git a/pgml-extension/requirements.txt b/pgml-extension/requirements.txt
index 1884801e5..3fdfeb4b7 100644
--- a/pgml-extension/requirements.txt
+++ b/pgml-extension/requirements.txt
@@ -25,3 +25,4 @@ transformers==4.31.0
 xgboost==1.7.6
 langchain==0.0.237
 einops==0.6.1
+pynvml==11.5.0
diff --git a/pgml-extension/src/bindings/transformers/mod.rs b/pgml-extension/src/bindings/transformers/mod.rs
index 85f71d3c8..d94e87de7 100644
--- a/pgml-extension/src/bindings/transformers/mod.rs
+++ b/pgml-extension/src/bindings/transformers/mod.rs
@@ -348,6 +348,8 @@ pub fn load_dataset(
 }
 
 pub fn clear_gpu_cache(memory_usage: Option<f32>) -> Result<bool> {
+    crate::bindings::venv::activate();
+
     Python::with_gil(|py| -> Result<bool> {
         let clear_gpu_cache: Py<PyAny> = PY_MODULE.getattr(py, "clear_gpu_cache")?;
         let success = clear_gpu_cache