diff --git a/1/model.py b/1/model.py
index 933f02d..bb23a16 100644
--- a/1/model.py
+++ b/1/model.py
@@ -3,6 +3,8 @@ 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):