Tests: BatchNormThroughTime

[jeshraghian]

Add new tests for BNTT1d and 2d.

[saiftyfirst]

Would do this.

Does this look like how this was meant to be used (Network from Tutorial 5 modified to use BatchNorm. Modifications have a
have a trailing comment - "# Batch_Normalization")

class Net(nn.Module):
    def __init__(self):
        super().__init__()

        # Initialize layers
        self.fc1 = nn.Linear(num_inputs, num_hidden)
        self.bn1d = snn.BatchNormTT1d(num_hidden, num_steps) # Batch_Normalization
        self.lif1 = snn.Leaky(beta=beta)
        self.fc2 = nn.Linear(num_hidden, num_outputs)
        self.lif2 = snn.Leaky(beta=beta)

    def forward(self, x):

        # Initialize hidden states at t=0
        mem1 = self.lif1.init_leaky()
        mem2 = self.lif2.init_leaky()
        
        # Record the final layer
        spk2_rec = []
        mem2_rec = []

        for step in range(num_steps):
            cur1 = self.fc1(x)
            curr1_bn = self.bn1d[step](curr1) # Batch_Normalization 
            spk1, mem1 = self.lif1(curr1_bn, mem1)
            cur2 = self.fc2(spk1)
            spk2, mem2 = self.lif2(cur2, mem2)
            spk2_rec.append(spk2)
            mem2_rec.append(mem2)

        return torch.stack(spk2_rec, dim=0), torch.stack(mem2_rec, dim=0)
        
# Load the network onto CUDA if available
net = Net().to(device)`