Implementing spatio-temporal convolutions in pytorch

I am trying to implement a layer to perform the (2+1)D convolutions described in this paper: https://arxiv.org/pdf/1711.11248.pdf

The basic idea is as follows: Let's say I have a 3D convolutional layer that takes in an input with $N_{i-1}$ channels, and performs $N_i$ 3D convolutions on a 3D tensor representing $T$ 2D frames taken over $T$ timesteps.

Instead of doing 3D convolutions over the tensor, the approach is to replace the 3D convolution with a 2D convolution followed by a 1D convolution along the temporal axis. In particular, if you wanted to perform $N_i$ 3D convolutions with kernels of size $N_{i-1} \times t \times d \times d$, you instead perform $M_i$ 2D convolutions with filters of size $N_{i-1} \times 1 \times d \times d$ followed by $N_i$ 1D convolutions along the temporal axis of size $M_i \times t \times 1 \times 1$.

Here is the tensorflow implementation of this type of layer: https://github.com/facebookresearch/R2Plus1D/blob/master/lib/models/video_model.py

I am having trouble implementing this type of layer in pytorch, however. This is what I have right now, and I'm clearly doing something wrong. I would appreciate it if someone could point me in the right direction as to how I would go about performing this type of convolution. I'm setting the intermediate channel dimension $M_i$ to 20 in this snippet:

class hybrid3d(nn.Module):

def __init__(self, n_classes):
super(hybrid3d, self).__init__()
t, d, d = 3, 3, 3
self.conv1_1 = nn.Conv2d(3, 20, kernel_size=(1, d, d), stride=[1,1,1] , padding=[0,0,0])
self.relu1_1 = nn.PReLU()
self.conv1_2 = nn.Conv1d(20, 12, kernel_size=(t, 1, 1), stride=[1,1,1], padding=[0,0,0])
self.relu1_2 = nn.PReLU()

def forward(self, x):
x = self.relu1_1(self.conv1_1(x))
x = self.relu1_2(self.conv1_2(x))
return x


import torch as t

class Net_1(t.nn.Module):
def __init__(self):
super(Net_1, self).__init__()
self.conv3d = t.nn.Conv3d(3, 8, (2, 3, 3))
self.relu = t.nn.ReLU()

def forward(self, x):
x = self.relu(self.conv3d(x))
return x

net_1 = Net_1()

video_1 = t.empty(1, 3, 6, 72, 108).normal_()

%%time
net_1(video_1).shape


Wall time: 18 ms

torch.Size([1, 8, 5, 70, 106])

class Net_2(t.nn.Module):
def __init__(self):
super(Net_2, self).__init__()
self.conv2d = t.nn.Conv2d(3, 8, (3, 3))
self.conv1d = t.nn.Conv1d(8, 8, (2))
self.flatten = t.nn.Flatten(start_dim=2)
self.relu = t.nn.ReLU()

def forward(self, x):
x = self.relu(self.conv2d(x))
x = self.flatten(x)
x.transpose_(0, 2)
x = self.relu(self.conv1d(x))
x.transpose_(0, 2)
x.transpose_(0, 1)
return x.view(5, 8, 70, 106)

net_2 = Net_2()

video_2 = t.empty(6, 3, 72, 108).normal_()

%%time
net_2(video_2).shape


Wall time: 58 ms

torch.Size([8, 5, 70, 106])