Feature/custom loaders (#526)

* Add with_loader

* Formatting

* Formatting

* Update changelog

* Remove sampler

* Update test pass

* Hide loader key

* Unhide loader key

* Revert to a5174b

* Add end to end

* Fix broken tests

* Update CHANGELOG.md

CHANGELOG.md

@@ -5,6 +5,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ## [Unreleased]
 ### Added
+- Added ``with_loader`` trial method that allows running of custom batch loaders
 ### Changed
 ### Deprecated
 ### Removed

tests/test_end_to_end.py

@@ -79,6 +79,27 @@ def test_basic_checkpoint(self):
         import os
         os.remove('test.pt')
 
+    def test_with_loader(self):
+        p = torch.tensor([2.0, 1.0, 10.0])
+        training_steps = 2
+
+        model = Net(p)
+        optim = torch.optim.SGD(model.parameters(), lr=0.01)
+        test_var = {'loaded': False}
+
+        def custom_loader(state):
+            state[torchbearer.X], state[torchbearer.Y_TRUE] = None, None
+            test_var['loaded'] = True
+
+        tbmodel = torchbearer.Trial(model, optim, loss, callbacks=[torchbearer.callbacks.MostRecent(filepath='test.pt')]).for_train_steps(training_steps).for_val_steps(1)
+        tbmodel.with_loader(custom_loader)
+        self.assertTrue(not test_var['loaded'])
+        tbmodel.run(1)
+        self.assertTrue(test_var['loaded'])
+
+        import os
+        os.remove('test.pt')
+
     def test_only_model(self):
         p = torch.tensor([2.0, 1.0, 10.0])