From 83c664e65a56eb353654b7c1536e7c59e8a7dcc4 Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Mon, 18 Apr 2022 17:51:39 -0700
Subject: [PATCH 1/2] Updated mlp tests

---
 mattport/nerf/field_modules/mlp.py            |  2 +-
 .../nerf/field_modules/test_field_outputs.py  |  8 ++---
 tests/nerf/field_modules/test_mlp.py          | 32 +++++++++++++++++++
 3 files changed, 37 insertions(+), 5 deletions(-)
 create mode 100644 tests/nerf/field_modules/test_mlp.py

diff --git a/mattport/nerf/field_modules/mlp.py b/mattport/nerf/field_modules/mlp.py
index 57c99a900..25deadd01 100644
--- a/mattport/nerf/field_modules/mlp.py
+++ b/mattport/nerf/field_modules/mlp.py
@@ -17,7 +17,7 @@ class MLP(FieldModule):
         out_dim: int,
         num_layers: int,
         layer_width: int,
-        activation: Optional[nn.Module] = None,
+        activation: Optional[nn.Module] = nn.ReLU(),
         out_activation: Optional[nn.Module] = None,
     ) -> None:
         """Initialize parameters of mulilayer perceptron.
diff --git a/tests/nerf/field_modules/test_field_outputs.py b/tests/nerf/field_modules/test_field_outputs.py
index cec54dd90..e7839ccfd 100644
--- a/tests/nerf/field_modules/test_field_outputs.py
+++ b/tests/nerf/field_modules/test_field_outputs.py
@@ -1,5 +1,5 @@
 """
-Render Head Tests
+Field output tests
 """
 import pytest
 import torch
@@ -9,7 +9,7 @@ from mattport.nerf.field_modules import field_outputs
 
 
 def test_field_output():
-    """Test rgb render head"""
+    """Test render output"""
     in_dim = 6
     out_dim = 4
     activation = nn.ReLU()
@@ -26,7 +26,7 @@ def test_field_output():
     assert y.shape[-1] == out_dim
 
 
-def test_density_head():
+def test_density_output():
     """Test rgb output"""
     in_dim = 6
     density_head = field_outputs.DensityFieldOutput(in_dim)
@@ -41,7 +41,7 @@ def test_density_head():
     assert y.shape[-1] == 1
 
 
-def test_rgb_head():
+def test_rgb_output():
     """Test rgb output"""
     in_dim = 6
     rgb_head = field_outputs.RGBFieldOutput(in_dim)
diff --git a/tests/nerf/field_modules/test_mlp.py b/tests/nerf/field_modules/test_mlp.py
new file mode 100644
index 000000000..498133883
--- /dev/null
+++ b/tests/nerf/field_modules/test_mlp.py
@@ -0,0 +1,32 @@
+"""
+MLP Test
+"""
+import pytest
+import torch
+from torch import nn
+
+from mattport.nerf.field_modules import MLP
+
+
+def test_mlp():
+    """Test mlp"""
+    in_dim = 6
+    out_dim = 10
+    num_layers = 2
+    layer_width = 32
+    out_activation = nn.ReLU()
+    mlp = MLP(
+        in_dim=in_dim, out_dim=out_dim, num_layers=num_layers, layer_width=layer_width, out_activation=out_activation
+    )
+    assert mlp.get_out_dim() == out_dim
+
+    x = torch.ones((9, in_dim))
+
+    mlp.build_nn_modules()
+    y = mlp(x)
+
+    assert y.shape[-1] == out_dim
+
+
+if __name__ == "__main__":
+    test_mlp()

From aa2525e5cbbeed3d537000156d86677c8b3c9e0b Mon Sep 17 00:00:00 2001
From: Matt Tancik <tancik@berkeley.edu>
Date: Mon, 18 Apr 2022 17:54:29 -0700
Subject: [PATCH 2/2] Update imports

---
 tests/nerf/field_modules/test_mlp.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/nerf/field_modules/test_mlp.py b/tests/nerf/field_modules/test_mlp.py
index 498133883..c25c29528 100644
--- a/tests/nerf/field_modules/test_mlp.py
+++ b/tests/nerf/field_modules/test_mlp.py
@@ -1,7 +1,6 @@
 """
 MLP Test
 """
-import pytest
 import torch
 from torch import nn