fix errors in broadcasting

    def forward(self, src):

        """Forward pass."""

        if self.to_log:

            # find the batch size

            batch_size = src.shape[0] // self.weights.shape[0]

            # repeat the weights

            weights = self.weights.repeat(batch_size)

            # multiply the weights

            src = src * weights

            # multiply the weights, use tensor broadcasting to deal with batch

            src = src * self.weights

            # take the log

            src = torch.log(src)