Suppose I have a Module:
Class A(nn.Module):
...
def __init__(...)
    self.word_embedding = nn.Embedding(size,hidden)
...

Can I use self.word_embedding outside of the class A forward function? Will this cause any problems with gradients?
For instance,

if __name__ == '__main__':
   ...
X = torch.ones(...)
model = A(...)
output = model(...)
X = model.word_embedding(X)
....

I know I can extract the embedding to a different module, but I want to keep A intact, with most of the job done in A’s forward.
Thanks!