Update app.py

app.py

@@ -12,7 +12,7 @@ os.environ['CUDA_LAUNCH_BLOCKING'] = '1'
 
 args = None
 lm_generation = None
-
+print(torch.__version__)
 
 def init_args():
     global args
@@ -43,11 +43,10 @@ def init_model():
     torch.set_default_tensor_type(torch.FloatTensor)
     model = load_model(model, args.load_model_path)
     model.eval()
-    print(model)
 
-    print(torch.cuda.max_memory_allocated() / 1024 ** 3)
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     model.to(device)
+    print(torch.cuda.max_memory_allocated() / 1024 ** 3)
     lm_generation = LmGeneration(model, args.tokenizer)