This has to be done when I define the ShiftLayer right?
class ShiftLayer(nn.Module):
    def __init__(self, b):
        super().__init__()
        self.register_buffer('b', b)
        self.b = nn.Parameter(torch.tensor(...))
        
    def forward(self, x):
        return x + self.b

like this?