nn.Conv1d consumes too much GPU memory - pytorch

Question

I have a ResBlock as below, which can change feature vector length as well. However, It consumes way too much GPU memory. In fact, one ResBlock like this alone can consume as much as 2.3GB of GPU memory, which causes CUDA_OUT_OF_MEMORY all the time.
Typical input size(batch size included): (65536, 256) or (65536, 63)
Typical output size(per ResBlock): (65536, 256)
The UpwardsConv1d module can change feature vector length with convolution.
You might think the batch size is too big, but a linear layer can handle that very well, which only consumes around 100MB of GPU memory per layer. There's no way that the Conv1d layer can't handle that with significantly fewer trainable parameters.
nn.Linear(63, 256) trainable parameters: 16384
ResBlock(63, 256) trainable parameters: 104
class UpwardsConv1d(nn.Module):
"""
Increase feature vector length by flattening dimentions.
"""
def __init__(self, size_in, size_out=None, size_h=8, k_size=3):
super().__init__()
self.size_in = size_in
if size_out is None:
self.size_out = size_in
else:
self.size_out = size_out
self.size_h = size_h
self.k_size = k_size
factor = math.ceil(self.size_out/self.size_in)
self.conv0 = nn.Conv1d(self.size_h, factor,
self.k_size, padding=(self.k_size-1)//2)
self.f = nn.Flatten(1, 2)
self.conv1 = nn.Conv1d(1, 1, self.size_in*factor + 1 - self.size_out)
def forward(self, x):
x = self.conv0(x)
x = self.f(x)
x = x.unsqueeze(-2)
x = self.conv1(x)
return x
class ResBlock(nn.Module):
"""
Act pretty much like a nn.Linear module, but uses convolution so have fewer trainable parameters.
"""
def __init__(self, size_in, size_out=None, size_h=2, k_size=3):
super().__init__()
self.size_in = size_in
if size_out is None:
self.size_out = size_in
else:
self.size_out = size_out
self.size_h = size_h
self.k_size = k_size
self.conv0 = nn.Conv1d(1, self.size_h, self.k_size,
padding=(self.k_size-1)//2)
if self.size_in == self.size_out:
self.conv1 = nn.Conv1d(self.size_h, 1, (self.k_size-1)//2)
else:
self.conv1 = UpwardsConv1d(
self.size_in, self.size_out, self.size_h, self.k_size)
def forward(self, x):
x = x.unsqueeze(-2)
x = self.conv0(x)
x = self.conv1(x)
x = x.squeeze(-2)
return x