Added a way to detach models and create them standalone
This commit is contained in:
@@ -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)
|
||||
)
|
||||
Reference in New Issue
Block a user