diff --git a/1/model.py b/1/model.py
index 933f02d..5cc83f2 100644
--- a/1/model.py
+++ b/1/model.py
@@ -3,6 +3,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 from peft import PeftModel, PeftConfig
 import numpy as np
 import json
+import torch
 
 class TritonPythonModel:
     def initialize(self, args):