I am training a Trigram model where I have created two dictionaries as below, to give index to each character.

```
stoi = {s: i+1 for i, s in enumerate(chars)}
stoi['.'] = 0
itos = {i:s for s,i in stoi.items()}
```

I have torch array of 3D. And I converted all the entries of array to probabilities as given below

```
P = (N+2).float()
P = P / P.sum(1, keepdims = True)
```

I didn’t voilate any broadcasting rules and the shape of P is 3D.

Now I am using torch.multinomial to draw a sample from my probablity distribution using the code below.

```
g = torch.Generator().manual_seed(2147483647)
for i in range(10):
out = []
ix = 0
while True:
p = P[ix]
ix = torch.multinomial(p, num_samples=1, replacement=True, generator=g)
out.append(itos[ix])
if ix == 0:
break
print(''.join(out))
```

But the above code does not draw samples from probability distribution, and gives me the below output.

""KeyError Traceback (most recent call last)

in <cell line: 3>()

7 p = P[ix]

8 ix = torch.multinomial(p, num_samples=1, replacement=True, generator=g)

----> 9 out.append(itos[ix])

10 if ix == 0:

11 break

KeyError: tensor([[10],

[14],

[18],

[10],

[ 1],

[14],

[ 0],

[17],

[ 5],

[15],

[ 0],

[25],

[ 0],

[ 3],

[ 1],

[ 2],

[ 5],

[16],

[ 5],

[25],

[ 8],

[ 0],

[17],

[ 8],

[24],

[ 1],

[25]]) “”

Can someone explain where there is an error in indexing, why itos dictionary doesn’t get proper index from probability distribution.