I customized a criterion function as my loss, while it turned out slower on GPU than running on CPU, in both calculating loss value and do loss.backward(). Here is my code defining the criterion.
import torch
import torch.nn as nn
from torch.distributions.negative_binomial import NegativeBinomial


class NegtiveBinomialLoss(nn.Module):
    def __init__(self):
        """
        this Module is designed as a customized loss criterion for negative binomial distribution,
        implementing -sum(log(likelihood))

        """
        super(NegtiveBinomialLoss, self).__init__()

    def forward(
        self,
        net_output: torch.autograd.Variable,
        y: torch.Tensor,
    ) -> torch.autograd.Variable:
        """
        Args:
            net_output: predict values, aka network output
            y: targets

        Returns: loss function variable

        """
        # alpha is number of success, lambda is number of failure
        param_alpha = net_output.t()[0]
        param_lambda = net_output.t()[1]
        log_likelihoods = torch.Tensor(len(y)).cuda()
        for i in range(len(y)):
            # In torch.distributions.negative_binomial.NegativeBinomial(total_count, probs),
            # total_count = success number, probs = failure prob
            nb = NegativeBinomial(total_count=param_alpha[i], probs=param_lambda[i] / (param_alpha[i] + param_lambda[i]))
            log_likelihoods[i] = nb.log_prob(y[i])
        return -log_likelihoods.sum()

I tried to put the criterion on cuda as well, but nothing changed.
criterion = NegtiveBinomialLoss().cuda()
loss = criterion(output, targets)
loss.backward()

Is there any method to speed it up?