File size: 239 Bytes
31726e5
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
"""Helper modules to build our networks."""
import torch as th


class Flatten(th.nn.Module):
    def __init__(self):
        super(Flatten, self).__init__()

    def forward(self, x):
        bs = x.shape[0]
        return x.view(bs, -1)