blog_10.py

class DNAConvNet(LightningModule):
    ...
   def _convert_to_jittable(self, module):
        for key, m in module._modules.items():
            if isinstance(m, MessagePassing) and m.jittable is not None:
                # Pytorch Geometric MessagePassing implements a `.jittable` function
                # which converts the current module into its jittable version.
                module._modules[key] = m.jittable()
            else:
                self._convert_to_jittable(m)
        return module

    def jittable(self):
        for key, m in self._modules.items():
            self._modules[key] = self._convert_to_jittable(m)