diff --git a/hgt.py b/hgt.py
index 3b8eb52..a722174 100755
--- a/hgt.py
+++ b/hgt.py
@@ -701,6 +701,7 @@ def add_budget(te, target_id, target_time, layer_data, budget):
         if (args.loss == "cross"):
             from torch_geometric.utils import remove_self_loops
             from torch_geometric.utils import add_self_loops
+            from torch_geometric.utils import negative_sampling
 
             EPS = 1e-15
             value = (node_rep[edge_index[0]] *