diff --git a/bird_cloud_gnn/callback.py b/bird_cloud_gnn/callback.py
index f4dc5d5..c55c180 100644
--- a/bird_cloud_gnn/callback.py
+++ b/bird_cloud_gnn/callback.py
@@ -5,8 +5,8 @@
 class TensorboardCallback:
     """Callback to populate Tensorboard"""
 
-    def __init__(self):
-        self.writer = SummaryWriter()
+    def __init__(self, **kwargs):
+        self.writer = SummaryWriter(**kwargs)
 
     def __call__(self, epoch_values):
         epoch = epoch_values["epoch"]