misc proposal network modifications (#578)

* misc improvements

* cleanup

* remove distortion

* dist loss

* only dist changed now

* changes

* remove comment
This commit is contained in:
Ethan Weber authored and GitHub committed 2022-09-29 16:31:22 -07:00
1 parent 5bc011e5fa
commit 81fa5f9ba5
7 files changed
+100 -28

No files matched your search

+4
View File
@@ -41,6 +41,10 @@ class Embedding(FieldModule):
def build_nn_modules(self) -> None:
self.embedding = torch.nn.Embedding(self.in_dim, self.out_dim)
def mean(self, dim=0):
"""Return the mean of the embedding weights along a dim."""
return self.embedding.weight.mean(dim)
def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
"""Call forward