Global Pooling Layers not jittable()

[bs1119]

I am trying to convert my PyG model to a torchscript. But, it seems the global pooling layers do not support jittable().
Can anyone help on how to convert global pooling layers? Maybe provide an example on how to rewrite global_mean_pooling to support jittable()?

Thanks.

Here is the error that I ran into:

RuntimeError:
Tried to access nonexistent attribute or method 'jittable' of type 'Tensor'.:
File "", line 54

    c1 = global_mean_pool(x1, batch).jittable()
         ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    
    c2_sum = global_add_pool(x2, batch).jittable()

[rusty1s]

The jittable call is only needed for MessagePassing like modules. The global_mean_pool operator should be jittable by default (nothing to do in this case), see here.