Error: Expected more than 1 value per channel when training

I’m using try & except because it only fails on the last batch. Do I need to convert the bn back to training mode if it only fails on the last batch? Will this mess up back propagation?

    def __init__(self, inchan, outchan):
        super(SiameseFCLayer, self).__init__()
        self.fc = nn.Linear(inchan, outchan)
        self.bn = nn.BatchNorm1d(outchan)
    def forward(self, x):
        out = self.fc(x)
        out = nn.functional.sigmoid(out)
        try:
            self.bn(out)
        except:
            self.bn.eval()
            self.bn(out)
            self.bn.train()
        return out