There’s an implementation here without using .reinforce. (https://github.com/JamesChuanggg/pytorch-REINFORCE/blob/master/reinforce_discrete.py 179).
basically, what you wrote is mostly correct, however you need to sample the categorical output (i.e the softmax output from the final output layer) to obtain an action, then you use the action to index your categorical output, and take log, then you can multiply the logprob by reward and sum it up like you wrote. The code snippets I’m talking about above are below:
probs = self.model(Variable(state)) # run net, get categorical probs output
action = probs.multinomial().data # sample to get action index
prob = probs[:, action[0,0]].view(1, -1) # index probs with action selection
log_prob = prob.log() # compute log prob

ith_step_loss = -log_prob*reward

You see the above loss implemented in line 52. Execept in this case they’re summing up loss across the entire trajectory, whereas I’ve written it above for one step action/reward.