fix bug in lr

        # transfer 2D prediction tensor to 2D x n_classes(=2) 

        complimentary = np.ones(np.shape(batch_output))

        complimentary = complimentary - batch_output

        batch_output = np.squeeze(np.stack(arrays = [complimentary,

                                                     batch_output],

                                            axis = 2))

        batch_output = np.concatenate([complimentary, batch_output],

                                            axis = batch_output.ndim-1)

        # reshape to batch_size x n_tasks x ...

        if batch_output.ndim == 3:

          batch_output = batch_output.transpose((1, 0, 2))