I seemed to have fixed it by doing
# linear layers
def linear(n_in, n_out):
    return nn.Linear(n_in, n_out)

def linear_layer(n_in, n_out, dp=False, dropout_rate = 0.5, bn = False, zero_bn = False, act=True):
    dropout = nn.Dropout(p=dropout_rate)
    bn = nn.BatchNorm1d(n_out)
    nn.init.constant_(bn.weight, 0. if zero_bn else 1.)

    layers = [linear(n_in, n_out)]
    if bn:
        layers.append(bn)
    if dp:
        layers.append(dropout)
    if act: 
        layers.append(act_fn)

    return nn.Sequential(*layers)