add import torch

younesbelkada · younesbelkada · commit 80df231ec7e3 · 2022-06-10T15:41:13.000+02:00
diff --git a/evaluation/generation/generate.py b/evaluation/generation/generate.py
@@ -2,6 +2,7 @@
 import json
 import datetime
 
+import torch
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
 def get_args():
