File "D:\AICode\ST-SSL\model\aug.py", line 53, in aug_topology drop_prob = torch.softmax(sim_mx[edge_mask], dim=0) RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu)

Echo-Ji / ST-SSL

ST-SSL (STSSL): Spatio-Temporal Self-Supervised Learning for Traffic Flow Forecasting/Prediction

131 stars 24 forks source link

File "D:\AICode\ST-SSL\model\aug.py", line 53, in aug_topology drop_prob = torch.softmax(sim_mx[edge_mask], dim=0) RuntimeError: indices should be either on cpu or on the same device as the indexed tensor (cpu) #3

Open XiangMeng666 opened 7 months ago

XiangMeng666 commented 7 months ago

I encountered the following problem when running the program, I debugging found that the sim_mx is in cpu, edge_mask is in gpu, and then I change the edge_mask to CPU, and the program can run successfully, i want to ensure whether changing this place is right?

my change is : edge_mask = (input_graph > 0).tril(diagonal=-1).cpu()

Echo-Ji commented 7 months ago

Yes, moving edge_mask to CPU is correct.