UPDATE: add sweeping requirements for tune and add tests

environment.yml

@@ -27,3 +27,4 @@ dependencies:
       - h5py==2.10.0
       - nbgrader==0.6.1
       - tune-sklearn==0.1.0
+      - hyperopt==0.2.5

tests/test_pytorch_model.py

@@ -29,3 +29,20 @@ def test_fit():
     torch_model.build_model()
     torch_model.fit(X, y)
 
+
+def test_sweep():
+
+    torch_model.build_model()
+    params = {"lr": [0.01, 0.02], "module__num_units": [10, 50]}
+    torch_model.sweep(params=params, X=X, y=y, search_algorithm="hyperopt")
+
+    torch_model.sweep(params=params, X=X, y=y, search_algorithm="bayesian")
+
+
+def test_predictor():
+
+    torch_model.build_model()
+    torch_model.fit(X, y)
+    y_hat = torch_model.predict(X)
+
+    assert y_hat.shape[0] == y.shape[0]