Binary Classification Accuracy calculation error

[akozlu]

anaconda3/lib/python3.6/site-packages/torchsample/modules/module_trainer.py", line 277, in fit
    metrics_logs = self.metric_container(output_batch, target_batch)
  File "/Users/aysekozlu/anaconda3/lib/python3.6/site-packages/torchsample/metrics.py", line 31, in __call__
    metric) 
anaconda3/lib/python3.6/site-packages/torchsample/modules/module_trainer.py", line 669, in calculate_loss
    return loss_fn(output_batch, target_batch)
anaconda3/lib/python3.6/site-packages/torchsample/metrics.py", line 86, in __call__
    self.correct_count += y_pred_round.eq(y_true).float().sum().data[0]
RuntimeError: The size of tensor a (2) must match the size of tensor b (64) at non-singleton dimension 1`

I think this happens when y_pred has more than one classes (i.e. the neural net output has more than 1 class.) This might be solved easily by finding index of element with highest probability
y_pred = np.argmax(y_pred, axis=1) # get position of max probability
I overwrote the class in the following way and it works, just letting you know.

class Binary_Classification(torchsample.metrics.BinaryAccuracy):
    def __call__(self, y_pred, y_true):
        # print('buraya_girdim')
        y_pred = y_pred.data.numpy()  # Turn it into np arrays
        y_true = y_true.data.numpy()
        y_pred = np.argmax(y_pred, axis=1)  # get position of max probability

        y_pred = y_pred.reshape(len(y_pred), 1)
        y_true = y_true.reshape(len(y_true), 1)
        self.correct_count += np.count_nonzero(y_pred == y_true)
        self.total_count += len(y_pred)
        accuracy = 100. * float(self.correct_count) / float(self.total_count)
        return accuracy