Update audiocraft/models/loaders.py

audiocraft/models/loaders.py

@@ -117,6 +117,10 @@ def load_lm_model(file_or_url_or_id: tp.Union[Path, str], device='cpu', depth='f
     _delete_param(cfg, 'conditioners.args.merge_text_conditions_p')
     _delete_param(cfg, 'conditioners.args.drop_desc_p')
     model = builders.get_lm_model(cfg)
+    if depth=='bfloat16':
+        model = model.to(torch.bfloat16)
+    if depth=='float16':
+        model = model.to(torch.float16)
     model.load_state_dict(pkg['best_state'])
     model.eval()
     model.cfg = cfg