`compute_prediction()` function does not work for a batch size of 1

[sumanabasu]

The compute_prediction() function breaks here for a single gene expression vector i.e. dimension of genes is [1, 977]

This is because scaled dosage returns a tensor of size 0.

However, it works for batch size of 2 onwards.