Yabin Li
2024-01-17 dcc310e1a8ff75df5b8d091a9e26883b85b8afec
funasr/auto/auto_model.py
@@ -146,7 +146,7 @@
        device = kwargs.get("device", "cuda")
        if not torch.cuda.is_available() or kwargs.get("ngpu", 0):
            device = "cpu"
            # kwargs["batch_size"] = 1
            kwargs["batch_size"] = 1
        kwargs["device"] = device
        
        if kwargs.get("ncpu", None):
@@ -183,9 +183,11 @@
            logging.info(f"Loading pretrained params from {init_param}")
            load_pretrained_model(
                model=model,
                init_param=init_param,
                path=init_param,
                ignore_init_mismatch=kwargs.get("ignore_init_mismatch", False),
                oss_bucket=kwargs.get("oss_bucket", None),
                scope_map=kwargs.get("scope_map", None),
                excludes=kwargs.get("excludes", None),
            )
        
        return model, kwargs