huangfu/InteractionImitation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

bugfix in deepsets

Browse Source
This commit is contained in:
Johannes Fischer
2021-07-22 18:44:12 +02:00
parent 91d052445e
commit 1c22bd6111
4 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
src/nets/deepsets.py
View File

@@ -63,7 +63,7 @@ class DeepSetsModule(nn.Module):
# shape (B, max_nv)
notnan_mask = torch.all(~torch.isnan(x), dim=-1)
# create zero tensor of shape (B, max_nv, latent_dim) to store phi evaluations in
latent = torch.zeros([*x.shape[:-1], self.latent_dim])
latent = torch.zeros([*x.shape[:-1], self.latent_dim], dtype=x.dtype)
# evaluate phi for all not NaN entries
# x[batch_dynamic_mask] has shape (notnan_mask.sum(), input_dim)
latent[notnan_mask] = self.phi(x[notnan_mask])