cp524 commited on
Commit
1af81da
·
1 Parent(s): a229791

bump batch sizes

Browse files
Files changed (1) hide show
  1. src/smc/inference.py +2 -2
src/smc/inference.py CHANGED
@@ -114,9 +114,9 @@ def _get_batch_size_based_on_gpu_mem_smc_grad(device, phi):
114
  if total_mem_gb < 24:
115
  batch_p = 1
116
  elif total_mem_gb < 48:
117
- batch_p = 2
118
- elif total_mem_gb < 70:
119
  batch_p = 4
 
 
120
  else:
121
  batch_p = 8
122
  elif phi <= 4:
 
114
  if total_mem_gb < 24:
115
  batch_p = 1
116
  elif total_mem_gb < 48:
 
 
117
  batch_p = 4
118
+ elif total_mem_gb < 70:
119
+ batch_p = 7
120
  else:
121
  batch_p = 8
122
  elif phi <= 4: