Vansh Chugh commited on
Commit
feb804f
·
1 Parent(s): f34b253

revert dead-code edit in loaders.py, real fix is in loopgen.py

Browse files
Files changed (1) hide show
  1. src/audiocraft/models/loaders.py +3 -5
src/audiocraft/models/loaders.py CHANGED
@@ -136,11 +136,9 @@ def load_lm_model_magnet(file_or_url_or_id: tp.Union[Path, str], compression_mod
136
  cfg.transformer_lm.span_len = cfg.masking.span_len
137
 
138
  # MAGNeT models v1 support only xformers backend.
139
- # from src.audiocraft.modules.transformer import set_efficient_attention_backend
140
- # if cfg.transformer_lm.memory_efficient:
141
- # set_efficient_attention_backend("xformers")
142
- # xformers has no kernel for ZeroGPU's Blackwell hardware; falling back to plain attention.
143
- cfg.transformer_lm.memory_efficient = False
144
 
145
  model = builders.get_lm_model(cfg)
146
  model.load_state_dict(pkg['best_state'])
 
136
  cfg.transformer_lm.span_len = cfg.masking.span_len
137
 
138
  # MAGNeT models v1 support only xformers backend.
139
+ from src.audiocraft.modules.transformer import set_efficient_attention_backend
140
+ if cfg.transformer_lm.memory_efficient:
141
+ set_efficient_attention_backend("xformers")
 
 
142
 
143
  model = builders.get_lm_model(cfg)
144
  model.load_state_dict(pkg['best_state'])