import torch.nn.functional as F

# Output for discrete actions
output = F.softmax(self.fc3(x), dim=-1)