hysts HF staff commited on
Commit
af9eb78
1 Parent(s): e080c62

Reduce max batch size for DSR

Browse files
Files changed (1) hide show
  1. model.py +2 -0
model.py CHANGED
@@ -241,6 +241,8 @@ class Model:
241
  start = time.perf_counter()
242
 
243
  srg = None if self.args.only_first_stage else SRGroup(self.args)
 
 
244
 
245
  elapsed = time.perf_counter() - start
246
  logger.info(f'--- done ({elapsed=:.3f}) ---')
 
241
  start = time.perf_counter()
242
 
243
  srg = None if self.args.only_first_stage else SRGroup(self.args)
244
+ if srg is not None:
245
+ srg.dsr.max_bz = 2
246
 
247
  elapsed = time.perf_counter() - start
248
  logger.info(f'--- done ({elapsed=:.3f}) ---')