fix useless code in memroy select

examples/test_all.sh

@@ -6,7 +6,7 @@ addr="192.168.1.107"
 partition_params=("ours" )
 partition_params=("ours" )
 #"metis" "ldg" "random")
 #"metis" "ldg" "random")
 #("ours" "metis" "ldg" "random")
 #("ours" "metis" "ldg" "random")
-partitions="12"
+partitions="8"
 node_per="4"
 node_per="4"
 nnodes="2"
 nnodes="2"
 node_rank="0"
 node_rank="0"

starrygl/module/memorys.py

@@ -484,7 +484,7 @@ class AsyncMemeoryUpdater(torch.nn.Module):
                     #print(self.gamma)
                     #print(self.gamma)
                     #print('tran {} {} {}\n'.format(transition_dense.max().item(),upd0[mask].max().item(),b.srcdata['his_mem'][mask].max().item()))
                     #print('tran {} {} {}\n'.format(transition_dense.max().item(),upd0[mask].max().item(),b.srcdata['his_mem'][mask].max().item()))
                 else:
                 else:
-                    upd0[mask] = b.srcdata['his_mem'][mask]
+                    upd0[mask] = updated_memory0[mask]
             #upd0[mask] = self.ceil_updater(his_mem, b.srcdata['his_mem'][mask])
             #upd0[mask] = self.ceil_updater(his_mem, b.srcdata['his_mem'][mask])
             #updated_memory = torch.where(mask.unsqueeze(1),self.gamma*updated_memory0 + (1-self.gamma)*(b.srcdata['his_mem'])
             #updated_memory = torch.where(mask.unsqueeze(1),self.gamma*updated_memory0 + (1-self.gamma)*(b.srcdata['his_mem'])
             # ,updated_memory0)
             # ,updated_memory0)