Computational graph

Say I have a PyTorch module and I want to obtain the corresponding unit-wise computational graph (forward pass). Is there way to do that in 2020? Many thanks in advance.