Update train_dec.py

ethanfetaya · Nov 19, 2019 · 415ef47 · 415ef47
1 parent 6a461ed
commit 415ef47
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/train_dec.py b/train_dec.py
@@ -95,8 +95,8 @@
 # Generate fully-connected interaction graph (sparse graphs would also work)
 off_diag = np.ones([args.num_atoms, args.num_atoms]) - np.eye(args.num_atoms)
 
-rel_rec = np.array(encode_onehot(np.where(off_diag)[1]), dtype=np.float32)
-rel_send = np.array(encode_onehot(np.where(off_diag)[0]), dtype=np.float32)
+rel_rec = np.array(encode_onehot(np.where(off_diag)[0]), dtype=np.float32)
+rel_send = np.array(encode_onehot(np.where(off_diag)[1]), dtype=np.float32)
 rel_rec = torch.FloatTensor(rel_rec)
 rel_send = torch.FloatTensor(rel_send)