Lord-Raven commited on
Commit
6d2fab3
·
1 Parent(s): 573c10d

CPU fallback.

Browse files
Files changed (1) hide show
  1. app.py +6 -3
app.py CHANGED
@@ -52,9 +52,12 @@ def classify(data_string, request: gradio.Request):
52
  # else:
53
  start_time = time.time()
54
  result = {}
55
- if ('cpu' in data):
56
- result = zero_shot_classification_cpu(data)
57
- else:
 
 
 
58
  result = zero_shot_classification_gpu(data)
59
  print(f"Classification @ [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] took {time.time() - start_time}.")
60
  return json.dumps(result)
 
52
  # else:
53
  start_time = time.time()
54
  result = {}
55
+ try:
56
+ if 'cpu' not in data:
57
+ result = zero_shot_classification_gpu(data)
58
+ except Exception as e:
59
+ print(f"GPU classification failed: {e}\nFall back to CPU.")
60
+ if not result:
61
  result = zero_shot_classification_gpu(data)
62
  print(f"Classification @ [{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}] took {time.time() - start_time}.")
63
  return json.dumps(result)