yuvalkirstain commited on
Commit
78264e2
1 Parent(s): eb993f9
Files changed (1) hide show
  1. app.py +2 -2
app.py CHANGED
@@ -58,9 +58,9 @@ def predict(prompt, image_1, image_2):
58
  print(f"Starting prediction for prompt: {prompt}")
59
  start_time = time.time()
60
  probs = calc_probs(prompt, [image_1, image_2])
61
- print(f"Prediction: {probs:.3f} ({time.time() - start_time:.2f} seconds, ) ")
62
  if device == "cuda":
63
- print(f"GPU mem used: {round(torch.cuda.max_memory_allocated(device) / 1024 / 1024 / 1024, 2)} GB")
64
  return str(round(probs[0], 3)), str(round(probs[1], 3))
65
 
66
 
 
58
  print(f"Starting prediction for prompt: {prompt}")
59
  start_time = time.time()
60
  probs = calc_probs(prompt, [image_1, image_2])
61
+ print(f"Prediction: {probs} ({time.time() - start_time:.2f} seconds, ) ")
62
  if device == "cuda":
63
+ print(f"GPU mem used: {round(torch.cuda.max_memory_allocated(device) / 1024 / 1024 / 1024, 2)}/{round(torch.cuda.get_device_properties(device).total_memory / 1024 / 1024 / 1024, 2)} GB")
64
  return str(round(probs[0], 3)), str(round(probs[1], 3))
65
 
66