From 412d3ff57d01d7e8c0889f686e31836170c4bfe3 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Sat, 11 Nov 2023 01:00:43 -0500
Subject: [PATCH] Refactor.

---
 comfy/ops.py | 24 +++++++++---------------
 1 file changed, 9 insertions(+), 15 deletions(-)

diff --git a/comfy/ops.py b/comfy/ops.py
index 610d54584fa..0bfb698aa7f 100644
--- a/comfy/ops.py
+++ b/comfy/ops.py
@@ -1,29 +1,23 @@
 import torch
 from contextlib import contextmanager
 
-class Linear(torch.nn.Module):
-    def __init__(self, in_features: int, out_features: int, bias: bool = True,
-                 device=None, dtype=None) -> None:
-        factory_kwargs = {'device': device, 'dtype': dtype}
-        super().__init__()
-        self.in_features = in_features
-        self.out_features = out_features
-        self.weight = torch.nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs))
-        if bias:
-            self.bias = torch.nn.Parameter(torch.empty(out_features, **factory_kwargs))
-        else:
-            self.register_parameter('bias', None)
-
-    def forward(self, input):
-        return torch.nn.functional.linear(input, self.weight, self.bias)
+class Linear(torch.nn.Linear):
+    def reset_parameters(self):
+        return None
 
 class Conv2d(torch.nn.Conv2d):
     def reset_parameters(self):
         return None
 
+class Conv3d(torch.nn.Conv3d):
+    def reset_parameters(self):
+        return None
+
 def conv_nd(dims, *args, **kwargs):
     if dims == 2:
         return Conv2d(*args, **kwargs)
+    elif dims == 3:
+        return Conv3d(*args, **kwargs)
     else:
         raise ValueError(f"unsupported dimensions: {dims}")