Update modules/pmbl.py
Browse files- modules/pmbl.py +1 -1
modules/pmbl.py
CHANGED
@@ -152,7 +152,7 @@ class PMBL:
|
|
152 |
|
153 |
conn.close()
|
154 |
|
155 |
-
|
156 |
def generate_topic(self, prompt, response):
|
157 |
llm = Llama(model_path=self.model_path, n_ctx=1690, n_threads=8, n_gpu_layers=-1, mlock=True)
|
158 |
llm = llm.to("cuda") # Move the model to the GPU
|
|
|
152 |
|
153 |
conn.close()
|
154 |
|
155 |
+
@spaces.GPU
|
156 |
def generate_topic(self, prompt, response):
|
157 |
llm = Llama(model_path=self.model_path, n_ctx=1690, n_threads=8, n_gpu_layers=-1, mlock=True)
|
158 |
llm = llm.to("cuda") # Move the model to the GPU
|