fix minor imports

Author: ilyes319

botnet/modules/loss.py

@@ -1,5 +1,5 @@
 import torch
-from torch_geometric.data import Batch
+from botnet.tools.torch_geometric.batch import Batch
 
 from botnet.tools import TensorDict
 

botnet/tools/train.py

@@ -102,7 +102,7 @@ def train(
 def take_step(
     model: torch.nn.Module,
     loss_fn: torch.nn.Module,
-    batch: torch_geometric.data.Batch,
+    batch: torch_geometric.batch.Batch,
     optimizer: torch.optim.Optimizer,
     ema: Optional[ExponentialMovingAverage],
     device: torch.device,