An example:
import torch
import time
import numpy as np
from torch.nn import functional as F
from torch.distributions.multivariate_normal import MultivariateNormal

_eps = 1e-20

def log_sum_exp(tensor, keepdim=True):
    """
    Numerically stable implementation for the `LogSumExp` operation. The
    summing is done along the last dimension.
    """
    max_val = tensor.max(dim=-1, keepdim=True)[0]
    return max_val + (tensor - max_val).exp().sum(dim=-1, keepdim=keepdim).log()


def loglik_mixture(x, z, c):
    """
    Calculate the log likelihood of a mixture of distributions
    """
    # Reshape the data
    x_reshape = x.contiguous().view(-1, x.shape[-1]).unsqueeze(0)
    x_repeat = x_reshape.repeat(c.shape[0], 1, 1)
    
    # Reshape the mixture weights
    z_reshape = z.view(-1, c.shape[0]) + _eps
    
    # Observation distribution parameters
    mu = x.data.new_zeros(x.shape[-1])  # zero mean
    sigma = c.unsqueeze(1)

    # Likelihood of each observation
    pdf = MultivariateNormal(mu, sigma).log_prob(x_repeat).t()
    
    # Log sum exp trick to sum the log probs, reshape so back in batch, seq_length
    return log_sum_exp(pdf + z_reshape.log()).reshape(*z.shape[:-1])

n_dim = 8
n_states = 6
n_batch = 40
seq_length = 200

c = torch.tensor(np.array([np.eye(n_dim)/i for i in range(1,n_states+1)]), requires_grad=True).float()
z = torch.randn([n_batch, seq_length, n_states], requires_grad=True).float()
z = F.softmax(z, dim=-1)
x = torch.randn([n_batch, seq_length, n_dim], requires_grad=False).float()

start = time.time()
loss = -loglik_mixture(x, z, c).sum()
print('Forwards pass: ', time.time() - start)

start = time.time()
loss.backward()
print('Backwards pass: ', time.time() - start)

This function seems to be the greatest issue in my code.
In 0.4.0 the forwards pass takes 0.0167 seconds and the backwards pass 0.0160 seconds.
In 1.1.0, the forwards pass takes 0.9668 seconds and the backwards pass 287.109 seconds.
Attached below are the bottleneck reports from 0.4.0 and 1.1.0, respectively:
--------------------------------------------------------------------------------
  Environment Summary
--------------------------------------------------------------------------------
PyTorch 0.4.0 not compiled w/ CUDA
Running with Python 3.6 and 

`pip list` truncated output:
Unable to fetch
--------------------------------------------------------------------------------
  cProfile output
--------------------------------------------------------------------------------
         290113 function calls (285732 primitive calls) in 2.371 seconds

   Ordered by: internal time
   List reduced from 2464 to 15 due to restriction <15>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
    32/29    0.611    0.019    0.617    0.021 {built-in method _imp.create_dynamic}
      210    0.592    0.003    0.661    0.003 <frozen importlib._bootstrap_external>:830(get_data)
      981    0.147    0.000    0.147    0.000 {built-in method posix.stat}
        1    0.128    0.128    0.128    0.128 {method 'run_backward' of 'torch._C._EngineBase' objects}
       46    0.079    0.002    0.079    0.002 {built-in method io.open}
      210    0.069    0.000    0.069    0.000 {method 'read' of '_io.FileIO' objects}
        3    0.061    0.020    0.061    0.020 {built-in method numpy.core.multiarray.fromfile}
      210    0.048    0.000    0.048    0.000 {built-in method marshal.loads}
     1200    0.041    0.000    0.121    0.000 /Users/askates/anaconda3/envs/torch/lib/python3.6/site-packages/torch/nn/_functions/rnn.py:48(GRUCell)
     2400    0.039    0.000    0.039    0.000 {built-in method addmm}
        2    0.030    0.015    0.049    0.024 /Users/askates/anaconda3/envs/torch/lib/python3.6/site-packages/torch/distributions/multivariate_normal.py:67(_batch_mahalanobis)
  901/848    0.024    0.000    0.065    0.000 {built-in method builtins.__build_class__}
      357    0.020    0.000    0.020    0.000 {method 'astype' of 'numpy.ndarray' objects}
        1    0.016    0.016    0.400    0.400 /Users/askates/Documents/Projects/model/src/train.py:94(train)
   410/98    0.015    0.000    0.024    0.000 /Users/askates/anaconda3/envs/torch/lib/python3.6/sre_parse.py:470(_parse)


--------------------------------------------------------------------------------
  autograd profiler output (CPU mode)
--------------------------------------------------------------------------------
        top 15 events sorted by cpu_time_total

------------------  ---------------  ---------------  ---------------  ---------------  ---------------
Name                       CPU time        CUDA time            Calls        CPU total       CUDA total
------------------  ---------------  ---------------  ---------------  ---------------  ---------------
mul                     14247.185us          0.000us                1      14247.185us          0.000us
mul                     13269.355us          0.000us                1      13269.355us          0.000us
MulBackward1             8679.297us          0.000us                1       8679.297us          0.000us
mul                      8674.418us          0.000us                1       8674.418us          0.000us
sum                      5864.228us          0.000us                1       5864.228us          0.000us
_sum                     5862.115us          0.000us                1       5862.115us          0.000us
sum                      5621.976us          0.000us                1       5621.976us          0.000us
_sum                     5620.054us          0.000us                1       5620.054us          0.000us
ExpandBackward           3796.324us          0.000us                1       3796.324us          0.000us
sum                      3794.368us          0.000us                1       3794.368us          0.000us
_sum                     3792.730us          0.000us                1       3792.730us          0.000us
sub                      1926.867us          0.000us                1       1926.867us          0.000us
PowBackward0             1900.918us          0.000us                1       1900.918us          0.000us
sub                      1782.307us          0.000us                1       1782.307us          0.000us
repeat                   1180.078us          0.000us                1       1180.078us          0.000us

--------------------------------------------------------------------------------
  Environment Summary
--------------------------------------------------------------------------------
PyTorch 1.1.0a0+19addc7 compiled w/ CUDA 9.2
Running with Python 3.6 and CUDA 9.2.148

`pip3 list` truncated output:
numpy==1.14.4
torch==0.4.0
--------------------------------------------------------------------------------
  cProfile output
--------------------------------------------------------------------------------
         684208 function calls (678739 primitive calls) in 485.558 seconds

   Ordered by: internal time
   List reduced from 2460 to 15 due to restriction <15>

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1  479.227  479.227  479.227  479.227 {method 'run_backward' of 'torch._C._EngineBase' objects}
      210    1.310    0.006    1.359    0.006 <frozen importlib._bootstrap_external>:830(get_data)
   128000    1.246    0.000    1.246    0.000 {built-in method trtrs}
    35/34    0.958    0.027    0.959    0.028 {built-in method _imp.create_dynamic}
   256000    0.701    0.000    0.701    0.000 /Users/askates/anaconda3/envs/pytorch-1.0/lib/python3.6/site-packages/torch/tensor.py:434(<lambda>)
        1    0.339    0.339  482.546  482.546 /Users/askates/Documents/Projects/model/src/train.py:94(train)
        2    0.313    0.157    2.260    1.130 /Users/askates/anaconda3/envs/pytorch-1.0/lib/python3.6/site-packages/torch/distributions/multivariate_normal.py:30(<listcomp>)
        4    0.295    0.074    0.295    0.074 {built-in method stack}
        4    0.280    0.070    0.280    0.070 {built-in method gru}
      968    0.100    0.000    0.100    0.000 {built-in method posix.stat}
       39    0.089    0.002    0.089    0.002 {built-in method io.open}
        2    0.062    0.031    2.616    1.308 /Users/askates/anaconda3/envs/pytorch-1.0/lib/python3.6/site-packages/torch/distributions/multivariate_normal.py:23(_batch_trtrs_lower)
      210    0.057    0.000    0.057    0.000 {built-in method marshal.loads}
      210    0.049    0.000    0.049    0.000 {method 'read' of '_io.FileIO' objects}
       14    0.035    0.002    0.035    0.002 {method 'reshape' of 'torch._C._TensorBase' objects}


--------------------------------------------------------------------------------
  autograd profiler output (CPU mode)
--------------------------------------------------------------------------------
        top 15 events sorted by cpu_time_total

------------------  ---------------  ---------------  ---------------  ---------------  ---------------
Name                       CPU time        CUDA time            Calls        CPU total       CUDA total
------------------  ---------------  ---------------  ---------------  ---------------  ---------------
stack                  144361.000us          0.000us                1     144361.000us          0.000us
stack                  137428.000us          0.000us                1     137428.000us          0.000us
StackBackward           95785.000us          0.000us                1      95785.000us          0.000us
unbind                  93179.000us          0.000us                1      93179.000us          0.000us
gru                     48150.000us          0.000us                1      48150.000us          0.000us
gru                     42979.000us          0.000us                1      42979.000us          0.000us
gru                     29062.000us          0.000us                1      29062.000us          0.000us
gru                     24434.000us          0.000us                1      24434.000us          0.000us
reshape                 22573.000us          0.000us                1      22573.000us          0.000us
clone                   22552.000us          0.000us                1      22552.000us          0.000us
reshape                 17592.000us          0.000us                1      17592.000us          0.000us
clone                   17574.000us          0.000us                1      17574.000us          0.000us
add                     17236.000us          0.000us                1      17236.000us          0.000us
SelectBackward          15354.000us          0.000us                1      15354.000us          0.000us
zeros                   15124.000us          0.000us                1      15124.000us          0.000us

--------------------------------------------------------------------------------
  autograd profiler output (CUDA mode)
--------------------------------------------------------------------------------
        top 15 events sorted by cpu_time_total

        Because the autograd profiler uses the CUDA event API,
        the CUDA time column reports approximately max(cuda_time, cpu_time).
        Please ignore this output if your code does not use CUDA.

-----------------  ---------------  ---------------  ---------------  ---------------  ---------------
Name                      CPU time        CUDA time            Calls        CPU total       CUDA total
-----------------  ---------------  ---------------  ---------------  ---------------  ---------------
stack                 221141.000us     221184.000us                1     221141.000us     221184.000us
stack                 149960.000us     150242.000us                1     149960.000us     150242.000us
gru                   143407.000us     143360.000us                1     143407.000us     143360.000us
StackBackward          95310.000us      95320.750us                1      95310.000us      95320.750us
unbind                 92935.000us      93348.500us                1      92935.000us      93348.500us
gru                    87878.000us      87873.500us                1      87878.000us      87873.500us
gru                    77095.000us      77952.000us                1      77095.000us      77952.000us
gru                    43305.000us      43295.469us                1      43305.000us      43295.469us
add                    34832.000us      34960.000us                1      34832.000us      34960.000us
reshape                31923.000us      31936.000us                1      31923.000us      31936.000us
clone                  31876.000us      31936.000us                1      31876.000us      31936.000us
add                    29661.000us      29696.000us                1      29661.000us      29696.000us
add                    22455.000us      22528.000us                1      22455.000us      22528.000us
add                    21128.000us      21248.000us                1      21128.000us      21248.000us
add                    20329.000us      23200.000us                1      20329.000us      23200.000us