Code Problem in permutations for symmetric hydrogens

[sunyuancheng]

Hi, thanks for the insight of this great work and for releasing the code!
But when reproducing training, I have encountered the following errors:

1. in model/model.py, in GeoMol, assign_neighobours, i got

File "/home/dgxtest/3D-pretrain/GeoMol-main/model/model.py", line 180, in assign_neighborhoods
RuntimeError: "mul_cuda" not implemented for 'Bool'
self.leaf_hydrogens[a] = self.leaf_hydrogens[a] * True if self.leaf_hydrogens[a].sum() > 1 else self.leaf_hydrogens[a] * False

I can see that this code is intended to executing a XNOR operation (not so convincing now due to error2), so I changed the logic into the following and fix the error

self.leaf_hydrogens[a] = ~(self.leaf_hydrogens[a] ^ True) if self.leaf_hydrogens[a].sum() > 1 else ~(self.leaf_hydrogens[a] ^ False)

2. But the following error ensues

File "/home/dgxtest/3D-pretrain/GeoMol-main/model/model.py", line 332, in ground_truth_local_stats
n_perms[0:len(perms), self.leaf_hydrogens[a]] = perms
'RuntimeError: shape mismatch: value tensor of shape [24, 4] cannot be broadcast to indexing result of shape [6, 4]'

in this case, self.leaf_hydrogens[a] is [True, True, True, True], thus leading to a permutation of length 24 in "perms" while "n_perms" is hardcoded in shape [6, 4]
I am not sure whether my modification in error1 leads to a wrong self.leaf_hydrogens in error2, would you please help me point it out? very much appreciated.

btw, I am using torch1.7.0+cu110 and torch-geometric 1.6.3 as metioned in issue #2.

[sunyuancheng]

Hi, I figured out that in tensor False * False = False, I rewrite it with & and fix the error, close the issue now.