Using torch.nn.ReLU instead of fvdb.nn.ReLU; some import cleanups

Signed-off-by: Jonathan Swartz <jonathan@jswartz.info>

Author: swahtz

fvdb/nn/simple_unet.py

@@ -30,23 +30,13 @@
 such as 3D semantic segmentation, shape completion, or volumetric reconstruction.
 """
 
-import math
-from typing import Any, Sequence
 
 import fvdb.nn as fvnn
 import torch
 import torch.nn as nn
-from fvdb.types import (
-    NumericMaxRank1,
-    NumericMaxRank2,
-    ValueConstraint,
-    to_Vec3i,
-    to_Vec3iBroadcastable,
-)
-from torch.profiler import record_function
-
-import fvdb
-from fvdb import ConvolutionPlan, Grid, GridBatch, JaggedTensor
+from fvdb.types import NumericMaxRank1
+
+from fvdb import ConvolutionPlan, GridBatch, JaggedTensor
 
 from .modules import fvnn_module
 
@@ -87,7 +77,7 @@ def __init__(
 
         self.conv = fvnn.SparseConv3d(in_channels, out_channels, kernel_size=kernel_size, stride=1, bias=False)
         self.batch_norm = fvnn.BatchNorm(out_channels, momentum=momentum)
-        self.relu = fvnn.ReLU(inplace=True)
+        self.relu = torch.nn.ReLU(inplace=True)
 
     def extra_repr(self) -> str:
         return (
@@ -107,7 +97,7 @@ def forward(
         x = self.conv(data, plan)
         out_grid = plan.target_grid_batch
         x = self.batch_norm(x, out_grid)
-        x = self.relu(x, out_grid)
+        x = self.relu(x)
         return x
 
 
@@ -170,7 +160,7 @@ def __init__(
 
         self.blocks = nn.ModuleList(layers)
 
-        self.final_relu = fvnn.ReLU(inplace=True)
+        self.final_relu = torch.nn.ReLU(inplace=True)
 
     def extra_repr(self) -> str:
         return (
@@ -199,7 +189,7 @@ def forward(
             data = block(data, plan)
 
         data = data + residual
-        data = self.final_relu(data, plan.target_grid_batch)
+        data = self.final_relu(data)
 
         return data