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)
|