Revert the config dicts

This commit is contained in:
henryruhs
2025-03-11 14:43:10 +01:00
parent 1dfd230fc5
commit 8101b15e1c
8 changed files with 119 additions and 170 deletions
+1 -2
View File
@@ -17,8 +17,7 @@ def export() -> None:
config_opset_version = CONFIG_PARSER.getint('exporting', 'opset_version')
os.makedirs(config_directory_path, exist_ok = True)
model = EmbeddingConverterTrainer.load_from_checkpoint(config_source_path, map_location = 'cpu')
model.eval()
model = EmbeddingConverterTrainer.load_from_checkpoint(config_source_path, map_location = 'cpu').eval()
model.ir_version = torch.tensor(config_ir_version)
input_tensor = torch.randn(1, 512)
torch.onnx.export(model, input_tensor, config_target_path, input_names = [ 'input' ], output_names = [ 'output' ], opset_version = config_opset_version)
+2 -2
View File
@@ -25,8 +25,8 @@ class EmbeddingConverterTrainer(LightningModule):
self.config_target_path = config_parser.get('training.model', 'target_path')
self.config_learning_rate = config_parser.getfloat('training.trainer', 'learning_rate')
self.embedding_converter = EmbeddingConverter()
self.source_embedder = torch.jit.load(self.config_source_path, map_location = 'cpu')
self.target_embedder = torch.jit.load(self.config_target_path, map_location = 'cpu')
self.source_embedder = torch.jit.load(self.config_source_path, map_location = 'cpu').eval()
self.target_embedder = torch.jit.load(self.config_target_path, map_location = 'cpu').eval()
self.mse_loss = nn.MSELoss()
def forward(self, source_embedding : Embedding) -> Embedding: