diff --git a/embedding_converter/src/exporting.py b/embedding_converter/src/exporting.py index e0fbae3..4ede026 100644 --- a/embedding_converter/src/exporting.py +++ b/embedding_converter/src/exporting.py @@ -17,7 +17,7 @@ def export() -> None: opset_version = CONFIG.getint('exporting', 'opset_version') makedirs(directory_path, exist_ok = True) - model = EmbeddingConverterTrainer.load_from_checkpoint(source_path) + model = EmbeddingConverterTrainer.load_from_checkpoint(source_path, map_location = 'cpu') model.eval() model.ir_version = torch.tensor(ir_version) input_tensor = (torch.randn(1, 512), ) diff --git a/face_swapper/src/exporting.py b/face_swapper/src/exporting.py index bae3287..6631fe1 100644 --- a/face_swapper/src/exporting.py +++ b/face_swapper/src/exporting.py @@ -17,7 +17,7 @@ def export() -> None: opset_version = CONFIG.getint('exporting', 'opset_version') makedirs(directory_path, exist_ok = True) - state_dict = torch.load(source_path).get('state_dict').get('generator') + state_dict = torch.load(source_path, map_location = 'cpu').get('state_dict').get('generator') model = Generator() model.load_state_dict(state_dict) model.eval()