Fix attribute loss

This commit is contained in:
henryruhs
2025-03-12 22:37:06 +01:00
parent 8f0ee4935b
commit 0d73bcf918
2 changed files with 8 additions and 4 deletions
+6 -3
View File
@@ -16,10 +16,13 @@ class Generator(nn.Module):
self.encoder.apply(init_weight)
self.generator.apply(init_weight)
def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tuple[Tensor, Tuple[Attribute, ...]]:
target_attributes = self.encoder(target_tensor)
def forward(self, source_embedding : Embedding, target_tensor : Tensor) -> Tensor:
target_attributes = self.encode_attributes(target_tensor)
output_tensor = self.generator(source_embedding, target_attributes)
return output_tensor, target_attributes
return output_tensor
def encode_attributes(self, input_tensor : Tensor) -> Tuple[Attribute, ...]:
return self.encoder(input_tensor)
def init_weight(module : nn.Module) -> None: