Update convolutions

hasktorch · Oct 27, 2020 · bcf0381 · bcf0381
1 parent 1b456b9
commit bcf0381
Showing 1 changed file with 11 additions and 3 deletions.
diff --git a/hasktorch/src/Torch/Typed/NN/Convolution.hs b/hasktorch/src/Torch/Typed/NN/Convolution.hs
@@ -99,7 +99,7 @@ instance
     Output
       (Conv1d inputChannelSize outputChannelSize kernelSize dtype device)
       (Tensor device dtype '[batchSize, inputChannelSize, inputSize], Proxy stride, Proxy padding) =
-      Tensor device dtype '[batchSize, outputChannelSize, outputSize]
+      Tensor device dtype '[batchSize, outputChannelSize, ConvOutputSize inputSize kernelSize stride padding]
   forward model (input, Proxy, Proxy) = conv1dForward @stride @padding model input
 
 instance
@@ -194,7 +194,7 @@ instance
     Output
       (Conv2d inputChannelSize outputChannelSize kernelSize0 kernelSize1 dtype device)
       (Tensor device dtype '[batchSize, inputChannelSize, inputSize0, inputSize1], Proxy stride, Proxy padding) =
-      Tensor device dtype '[batchSize, outputChannelSize, outputSize0, outputSize1]
+      Tensor device dtype '[batchSize, outputChannelSize, ConvOutputSize inputSize0 kernelSize0 (Torch.Typed.Aux.Fst stride) (Torch.Typed.Aux.Fst padding), ConvOutputSize inputSize1 kernelSize1 (Torch.Typed.Aux.Snd stride) (Torch.Typed.Aux.Snd padding)]
   forward model (input, Proxy, Proxy) = conv2dForward @stride @padding model input
 
 instance
@@ -296,7 +296,15 @@ instance
     Output
       (Conv3d inputChannelSize outputChannelSize kernelSize0 kernelSize1 kernelSize2 dtype device)
       (Tensor device dtype '[batchSize, inputChannelSize, inputSize0, inputSize1, inputSize2], Proxy stride, Proxy padding) =
-      Tensor device dtype '[batchSize, outputChannelSize, outputSize0, outputSize1, outputSize2]
+      Tensor
+        device
+        dtype
+        '[ batchSize,
+           outputChannelSize,
+           ConvOutputSize inputSize0 kernelSize0 (Fst3 stride) (Fst3 padding),
+           ConvOutputSize inputSize1 kernelSize1 (Snd3 stride) (Snd3 padding),
+           ConvOutputSize inputSize2 kernelSize2 (Trd3 stride) (Trd3 padding)
+         ]
   forward model (input, Proxy, Proxy) = conv3dForward @stride @padding model input
 
 instance