Update audiocraft/models/loaders.py

audiocraft/models/loaders.py

@@ -121,8 +121,10 @@ def load_lm_model(file_or_url_or_id: tp.Union[Path, str], device='cpu', depth='f
     model = builders.get_lm_model(cfg)
     if depth=='bfloat16':
         model = model.to(torch.bfloat16)
+        cfg.dtype = 'bfloat16'
     if depth=='float16':
         model = model.to(torch.float16)
+        cfg.dtype = 'float16'
     model.load_state_dict(pkg['best_state'])
     model.eval()
     model.cfg = cfg