diff --git a/setup.cfg b/setup.cfg
index e7ea71f..2f505fd 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -22,7 +22,6 @@ include_package_data = True
 install_requires =
     torch
     numpy
-    funlib.learn.torch @ git+https://github.com/funkelab/funlib.learn.torch.git
 
 [options.packages.find]
 where=src
diff --git a/src/leibnetz/leibnet.py b/src/leibnetz/leibnet.py
index de5d0da..bb4051c 100644
--- a/src/leibnetz/leibnet.py
+++ b/src/leibnetz/leibnet.py
@@ -394,7 +394,7 @@ def forward(self, inputs: dict[str, dict[str, Sequence[int | float]]]):
         # outputs is a dictionary of tensors
 
         # initialize buffer
-        self.buffer = inputs
+        self.buffer = {key: inputs[key] for key in self.input_keys}
 
         # march along nodes based on graph succession
         for flushable_list, node in zip(self.flushable_arrays, self.ordered_nodes):