probs = policy_network(state)
   # Note that this is equivalent to what used to be called multinomial
   m = Categorical(probs)
   action = m.sample()
   next_state, reward = env.step(action)
   loss = -m.log_prob(action) * reward
   loss.backward()
Pathwise derivative
The other way to implement these stochastic/policy gradients would be
to use the reparameterization trick from the "rsample()" method, where
the parameterized random variable can be constructed via a
parameterized deterministic function of a parameter-free random
variable. The reparameterized sample therefore becomes differentiable.
The code for implementing the pathwise derivative would be as follows:
params = policy_network(state)
   m = Normal(*params)
   # Any distribution with .has_rsample == True could work based on the application
   action = m.rsample()
   next_state, reward = env.step(action)  # Assuming that reward is differentiable
   loss = -reward
   loss.backward()
Distribution