Solved by InnovArul in post #2 
              
              
                How about this? 
import torch                                                           
from torch import nn                                                
from torch.distribitutions.bernoulli import Bernoulli

nwords, dim = 4, 3
emb = nn.Embedding(nwords, dim)
input = torch.LongTensor([[0, 1, 2, …