From 330fe2f0f4b1a6221ccca57aee89d972c2e9af22 Mon Sep 17 00:00:00 2001
From: suzgunmirac <msuzgun@college.harvard.edu>
Date: Fri, 1 Feb 2019 11:46:44 -0500
Subject: [PATCH] added more pooling layers

---
 namedtensor/torch_nn.py | 38 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 37 insertions(+), 1 deletion(-)

diff --git a/namedtensor/torch_nn.py b/namedtensor/torch_nn.py
index f548f7e..2b86949 100644
--- a/namedtensor/torch_nn.py
+++ b/namedtensor/torch_nn.py
@@ -95,6 +95,31 @@ class MaxPool3d(_Update, nn.MaxPool2d):
     pass
 
 
+class AvgPool1d (_Update, nn.AvgPool1d):
+    pass
+
+
+class AvgPool2d (_Update, nn.AvgPool2d):
+    pass
+
+
+class AvgPool3d (_Update, nn.AvgPool3d):
+    pass
+
+
+class MaxUnpool1d (_Update, nn.MaxUnpool1d):
+    pass
+
+
+class MaxUnpool2d (_Update, nn.MaxUnpool2d):
+    pass
+
+
+class MaxUnpool3d (_Update, nn.MaxUnpool3d):
+    pass
+
+
+
 _update = [
     "Linear",
     "Conv1d",
@@ -103,6 +128,12 @@ class MaxPool3d(_Update, nn.MaxPool2d):
     "MaxPool1d",
     "MaxPool2d",
     "MaxPool3d",
+    "AvgPool1d",
+    "AvgPool2d",
+    "AvgPool3d",
+    "MaxUnpool1d",
+    "MaxUnpool2d",
+    "MaxUnpool3d",
 ]
 
 
@@ -113,7 +144,12 @@ class MaxPool3d(_Update, nn.MaxPool2d):
 MaxPool1d.__doc__ = nn.MaxPool1d.__doc__
 MaxPool2d.__doc__ = nn.MaxPool2d.__doc__
 MaxPool3d.__doc__ = nn.MaxPool3d.__doc__
-
+AvgPool1d.__doc__ = nn.AvgPool1d.__doc__
+AvgPool2d.__doc__ = nn.AvgPool2d.__doc__
+AvgPool3d.__doc__ = nn.AvgPool3d.__doc__
+MaxUnpool1d.__doc__ = nn.MaxUnpool1d.__doc__
+MaxUnpool2d.__doc__ = nn.MaxUnpool2d.__doc__
+MaxUnpool3d.__doc__ = nn.MaxUnpool3d.__doc__
 
 class CrossEntropyLoss(_Loss, nn.CrossEntropyLoss):
     pass