diff --git a/densenet.py b/densenet.py
index 35321e3..52ce73b 100644
--- a/densenet.py
+++ b/densenet.py
@@ -15,6 +15,11 @@
 import sys
 import math
 
+## TODO: Implement Conv3x3 to follow Official PyTorch styling
+
+def flatten(x): 
+    return x.view(x.size(0), -1)
+
 class Bottleneck(nn.Module):
     def __init__(self, nChannels, growthRate):
         super(Bottleneck, self).__init__()
@@ -58,14 +63,14 @@ def forward(self, x):
 
 
 class DenseNet(nn.Module):
-    def __init__(self, growthRate, depth, reduction, nClasses, bottleneck):
+     def __init__(self, growthRate, depth, reduction, nClasses, bottleneck, compression):
         super(DenseNet, self).__init__()
-
+    
         nDenseBlocks = (depth-4) // 3
         if bottleneck:
             nDenseBlocks //= 2
 
-        nChannels = 2*growthRate
+        nChannels = 2 * growthRate if compression and bottleneck else 16 # They only do this for BC type
         self.conv1 = nn.Conv2d(3, nChannels, kernel_size=3, padding=1,
                                bias=False)
         self.dense1 = self._make_dense(nChannels, growthRate, nDenseBlocks, bottleneck)
@@ -84,6 +89,7 @@ def __init__(self, growthRate, depth, reduction, nClasses, bottleneck):
         nChannels += nDenseBlocks*growthRate
 
         self.bn1 = nn.BatchNorm2d(nChannels)
+        self.avgpool = nn.AvgPool2d(kernel_size=8)
         self.fc = nn.Linear(nChannels, nClasses)
 
         for m in self.modules():
@@ -105,12 +111,57 @@ def _make_dense(self, nChannels, growthRate, nDenseBlocks, bottleneck):
                 layers.append(SingleLayer(nChannels, growthRate))
             nChannels += growthRate
         return nn.Sequential(*layers)
-
-    def forward(self, x):
-        out = self.conv1(x)
-        out = self.trans1(self.dense1(out))
-        out = self.trans2(self.dense2(out))
-        out = self.dense3(out)
-        out = torch.squeeze(F.avg_pool2d(F.relu(self.bn1(out)), 8))
-        out = F.log_softmax(self.fc(out))
+    
+    def forward(self, x):                    # Output dims
+        out = self.conv1(x)                  # 32x32
+        
+        out = self.trans1(self.dense1(out))  # 16x16
+        out = self.trans2(self.dense2(out))  # 8x8
+        out = self.dense3(out)               # 8x8
+        
+        out = self.avgpool(out)              # 1x1
+        out = self.flatten(out)
+        out = self.fc(out) 
+        
         return out
+
+    
+def denseNet_40_12():
+    return DenseNet(12, 40, 1, 10, bottleneck=False, compression=False)
+
+def denseNet_100_12():
+    return DenseNet(12, 100, 1, 10, bottleneck=False, compression=False)
+
+def denseNet_100_24():
+    return DenseNet(24, 100, 1, 10, bottleneck=False, compression=False)
+
+def denseNetBC_100_12():
+    return DenseNet(12, 100, 0.5, 10, bottleneck=True, compression=True)
+
+def denseNetBC_250_24():
+    return DenseNet(24, 250, 0.5, 10, bottleneck=True, compression=True)
+
+def denseNetBC_190_40():
+    return DenseNet(40, 190, 0.5, 10, bottleneck=True, compression=True)
+
+    
+''' Did this little check:
+DenseNets implemented on the paper <https://arxiv.org/pdf/1608.06993.pdf>
+
++-------------+-------------+-------+--------------+
+|    Model    | Growth Rate | Depth | M. of Params |
++-------------+-------------+-------+--------------+
+|  DenseNet   |     12      |  40   |     1.02     |
++-------------+-------------+-------+--------------+
+|  DenseNet   |     12      |  100  |     6.98     |
++-------------+-------------+-------+--------------+
+|  DenseNet   |     24      |  100  |    27.249    |
++-------------+-------------+-------+--------------+
+| DenseNet-BC |     12      |  100  |    0.769     |
++-------------+-------------+-------+--------------+
+| DenseNet-BC |     24      |  250  |    15.324    |
++-------------+-------------+-------+--------------+
+| DenseNet-BC |     40      |  190  |    25.624    |
++-------------+-------------+-------+--------------+
+
+''''