Added a way to detach models and create them standalone

This commit is contained in:
Christian Risi
2025-10-10 18:43:20 +02:00
parent 15f203cad5
commit 92ae40013d
10 changed files with 164 additions and 20 deletions

View File

@@ -46,10 +46,17 @@ class TrainingModel(torch.nn.Module):
encoder_output, _ = self.__encoder((encoder_tensor, src_padding))
decoder_output, _, _, _, _ = self.__decoder(
(decoder_tensor, encoder_output, encoder_output, src_padding, tgt_padding)
decoder_output, _, _, _, _, _ = self.__decoder(
(decoder_tensor, encoder_output, encoder_output, src_padding, tgt_padding, False)
)
logits: torch.Tensor = self.__detokener(decoder_output)
return logits
def take_pieces(self):
return (
(self.__encoder_embedder, self.__encoder),
(self.__decoder_embedder, self.__decoder, self.__detokener)
)