OpenTransformer commited on
Commit
2d87d7b
·
verified ·
1 Parent(s): 7f8f08e

Make single-file dblock MoE self-contained

Browse files
Files changed (1) hide show
  1. agillm4_dblock_single_file.py +2 -2
agillm4_dblock_single_file.py CHANGED
@@ -3120,7 +3120,7 @@ def _train_phase(
3120
  break
3121
  if _profile_batch:
3122
  try:
3123
- import dblocks_train as _db_prof
3124
  _db_prof._profile_add(_DBS, "data_stream", time.perf_counter() - _data_t)
3125
  except Exception:
3126
  pass
@@ -3138,7 +3138,7 @@ def _train_phase(
3138
  except Exception:
3139
  pass
3140
  try:
3141
- import dblocks_train as _db_prof
3142
  _db_prof._profile_add(_DBS, "tensor", time.perf_counter() - _tensor_t)
3143
  except Exception:
3144
  pass
 
3120
  break
3121
  if _profile_batch:
3122
  try:
3123
+ _db_prof = sys.modules.get(__name__)
3124
  _db_prof._profile_add(_DBS, "data_stream", time.perf_counter() - _data_t)
3125
  except Exception:
3126
  pass
 
3138
  except Exception:
3139
  pass
3140
  try:
3141
+ _db_prof = sys.modules.get(__name__)
3142
  _db_prof._profile_add(_DBS, "tensor", time.perf_counter() - _tensor_t)
3143
  except Exception:
3144
  pass