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

changing default network, making deepsets network choose output dimension appropriately, making bc config to be called by ray

Browse Source
This commit is contained in:
Arec
2021-07-26 07:52:11 -07:00
parent 69359b5af3
commit d9436daaba
4 changed files with 42 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/nets/deepsets.py
View File

@@ -18,9 +18,9 @@ class DeepSetsModule(nn.Module):
super(DeepSetsModule, self).__init__()
self.input_dim = input_dim
self.latent_dim = latent_dim
self.output_dim = output_dim
self.phi = Phi(self.input_dim, phi_hidden_n, phi_hidden_dim, self.latent_dim)
self.rho = Phi(self.latent_dim, rho_hidden_n, rho_hidden_dim, self.output_dim)
self.rho = Phi(self.latent_dim, rho_hidden_n, rho_hidden_dim, output_dim)
self.output_dim = self.rho.output_dim
self.pooling = torch.sum
@staticmethod