PyTorch Geometric TorchScript generation with low level operations from aten library

[adhithadias]

Hi, I am trying to generate the TorchScript for the https://pytorch-geometric.readthedocs.io/en/latest/notes/jit.html simple graph in this document.

I am printing the model.graph and below is what I get but this code does not include any low-level operations used in the neural network graph, like the functions from aten::mm, aten::add

graph(%self : __torch__.models.GCNNet,
      %x.1 : Tensor,
      %edge_index.1 : Tensor):
  %27 : Function = prim::Constant[name="log_softmax"]()
  %25 : int = prim::Constant[value=3]()
  %16 : Function = prim::Constant[name="dropout"]()
  %10 : Function = prim::Constant[name="relu"]()
  %9 : bool = prim::Constant[value=0]()
  %6 : None = prim::Constant()
  %13 : float = prim::Constant[value=0.59999999999999998]() # /src/models.py:15:27
  %24 : int = prim::Constant[value=1]() # /src/models.py:17:36
  %3 : __torch__.GCNConvJittable_bce75b.GCNConvJittable_bce75b = prim::GetAttr[name="conv1"](%self)
  %x.3 : Tensor = prim::CallMethod[name="forward__0"](%3, %x.1, %edge_index.1, %6) # /src/models.py:13:12
  %x.5 : Tensor = prim::CallFunction(%10, %x.3, %9) # /src/models.py:14:12
  %14 : bool = prim::GetAttr[name="training"](%self)
  %x.7 : Tensor = prim::CallFunction(%16, %x.5, %13, %14, %9) # /src/models.py:15:12
  %18 : __torch__.GCNConvJittable_bce75b.___torch_mangle_0.GCNConvJittable_bce75b = prim::GetAttr[name="conv2"](%self)
  %x.9 : Tensor = prim::CallMethod[name="forward__0"](%18, %x.7, %edge_index.1, %6) # /src/models.py:16:12
  %28 : Tensor = prim::CallFunction(%27, %x.9, %24, %25, %6) # /src/models.py:17:15
  return (%28)

But it doesn't print out the low level operations used from the aten library?
Is there a way to generate those as well? Like the below torchscript which I got from a different non pytorch geometric code.

graph(%self : __torch__.LSTMCell,
 %input.1 : Tensor,
 %state.1 : (Tensor, Tensor)):
%23 : int = prim::Constant[value=4]() # ../eg.py:24:60
%24 : int = prim::Constant[value=1]() # ../eg.py:24:63
%hx.1 : Tensor, %cx.1 : Tensor = prim::TupleUnpack(%state.1)
%7 : Tensor = prim::GetAttr[name="weight_ih"](%self)
%8 : Tensor = **_aten::t_**(%7) # ../eg.py:22:33
%9 : Tensor = _**aten::mm**_(%input.1, %8) # ../eg.py:22:17
%10 : Tensor = prim::GetAttr[name="bias_ih"](%self)
%12 : Tensor = _**aten::add**_(%9, %10, %24) # ../eg.py:22:17
%14 : Tensor = prim::GetAttr[name="weight_hh"](%self)
%15 : Tensor = _**aten::t**_(%14) # ../eg.py:23:30
%16 : Tensor = _**aten::mm**_(%hx.1, %15) # ../eg.py:23:17
%18 : Tensor = _**aten::add**_(%12, %16, %24) # ../eg.py:22:17
%19 : Tensor = prim::GetAttr[name="bias_hh"](%self)
%gates.1 : Tensor = _**aten::add**_(%18, %19, %24) # ../eg.py:22:17
...

[rusty1s]

See #2127