John6666 commited on
Commit
88cd3e7
1 Parent(s): f788018

Upload utils.py

Browse files
Files changed (1) hide show
  1. utils.py +6 -1
utils.py CHANGED
@@ -8,7 +8,12 @@ def get_status(model_name: str):
8
  return client.get_model_status(model_name)
9
 
10
  def is_loadable(model_name: str, force_gpu: bool = False):
11
- status = get_status(model_name)
 
 
 
 
 
12
  gpu_state = isinstance(status.compute_type, dict) and "gpu" in status.compute_type.keys()
13
  if status is None or status.state not in ["Loadable", "Loaded"] or (force_gpu and not gpu_state):
14
  print(f"Couldn't load {model_name}. Model state:'{status.state}', GPU:{gpu_state}")
 
8
  return client.get_model_status(model_name)
9
 
10
  def is_loadable(model_name: str, force_gpu: bool = False):
11
+ try:
12
+ status = get_status(model_name)
13
+ except Exception as e:
14
+ print(e)
15
+ print(f"Couldn't load {model_name}.")
16
+ return False
17
  gpu_state = isinstance(status.compute_type, dict) and "gpu" in status.compute_type.keys()
18
  if status is None or status.state not in ["Loadable", "Loaded"] or (force_gpu and not gpu_state):
19
  print(f"Couldn't load {model_name}. Model state:'{status.state}', GPU:{gpu_state}")