File size: 261 Bytes
b05b302
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
import torch
import torch.nn as nn

class UncertaintyActionModel(nn.Module):
    def init(self):
        super().init()
        self.fc = nn.Linear(3, 3)  # confidence, risk, ambiguity

    def forward(self, x):
        return torch.softmax(self.fc(x), dim=-1)