ImportError: No module named 'torch_ext'

Hi, I was running the “quickstart” code on my win10. I used torch = 1.8.0 and python = 3.7 with cuda = 10.2.

The problems happened when I tried training the model in Jupyter:

optimizer = torch.optim.Adam(task.parameters(), lr=1e-3)
solver = core.Engine(task, train_set, valid_set, test_set, optimizer, gpus=[0],
                     batch_size=512) solver.train(num_epoch=100)

And this turned to:

---------------------------------------------------------------------------
ImportError                               Traceback (most recent call last)
~\AppData\Local\Temp/ipykernel_19744/1406504193.py in <module>
----> 1 solver.train(num_epoch=100)

d:\conda\envs\torchdrug\lib\site-packages\torchdrug-0.1.0-py3.7.egg\torchdrug\core\engine.py in train(self, num_epoch, batch_per_epoch)
    141                     batch = utils.cuda(batch, device=self.device)
    142 
--> 143                 loss, metric = model(batch)
    144                 if not loss.requires_grad:
    145                     raise RuntimeError("Loss doesn't require grad. Did you define any loss in the task?")

d:\conda\envs\torchdrug\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

d:\conda\envs\torchdrug\lib\site-packages\torchdrug-0.1.0-py3.7.egg\torchdrug\tasks\property_prediction.py in forward(self, batch)
     72         metric = {}
     73 
---> 74         pred = self.predict(batch, all_loss, metric)
     75 
     76         if all([t not in batch for t in self.task]):

d:\conda\envs\torchdrug\lib\site-packages\torchdrug-0.1.0-py3.7.egg\torchdrug\tasks\property_prediction.py in predict(self, batch, all_loss, metric)
    103     def predict(self, batch, all_loss=None, metric=None):
    104         graph = batch["graph"]
--> 105         output = self.model(graph, graph.node_feature.float(), all_loss=all_loss, metric=metric)
    106         pred = self.linear(output["graph_feature"])
    107         return pred

d:\conda\envs\torchdrug\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

d:\conda\envs\torchdrug\lib\site-packages\torchdrug-0.1.0-py3.7.egg\torchdrug\models\gin.py in forward(self, graph, input, all_loss, metric)
     74 
     75         for layer in self.layers:
---> 76             hidden = layer(graph, layer_input)
     77             if self.short_cut and hidden.shape == layer_input.shape:
     78                 hidden = hidden + layer_input

d:\conda\envs\torchdrug\lib\site-packages\torch\nn\modules\module.py in _call_impl(self, *input, **kwargs)
    887             result = self._slow_forward(*input, **kwargs)
    888         else:
--> 889             result = self.forward(*input, **kwargs)
    890         for hook in itertools.chain(
    891                 _global_forward_hooks.values(),

d:\conda\envs\torchdrug\lib\site-packages\torchdrug-0.1.0-py3.7.egg\torchdrug\layers\conv.py in forward(self, graph, input)
     89             update = checkpoint.checkpoint(self._message_and_aggregate, *graph.to_tensors(), input)
     90         else:
---> 91             update = self.message_and_aggregate(graph, input)
     92         output = self.combine(input, update)
     93         return output

d:\conda\envs\torchdrug\lib\site-packages\torchdrug-0.1.0-py3.7.egg\torchdrug\layers\conv.py in message_and_aggregate(self, graph, input)
    339     def message_and_aggregate(self, graph, input):
    340         adjacency = utils.sparse_coo_tensor(graph.edge_list.t()[:2], graph.edge_weight,
--> 341                                             (graph.num_node, graph.num_node))
    342         update = torch.sparse.mm(adjacency.t(), input)
    343         if self.edge_linear:

d:\conda\envs\torchdrug\lib\site-packages\torchdrug-0.1.0-py3.7.egg\torchdrug\utils\torch.py in sparse_coo_tensor(indices, values, size)
    160         size (list): size of the tensor
    161     """
--> 162     return torch_ext.sparse_coo_tensor_unsafe(indices, values, size)
    163 
    164 

d:\conda\envs\torchdrug\lib\site-packages\torchdrug-0.1.0-py3.7.egg\torchdrug\utils\torch.py in __getattr__(self, key)
     28             self.module = cpp_extension.load(self.name, self.sources, self.extra_cflags, self.extra_cuda_cflags,
     29                                              self.extra_ldflags, self.extra_include_paths, self.build_directory,
---> 30                                              self.verbose, **self.kwargs)
     31         return getattr(self.module, key)
     32 

d:\conda\envs\torchdrug\lib\site-packages\torch\utils\cpp_extension.py in load(name, sources, extra_cflags, extra_cuda_cflags, extra_ldflags, extra_include_paths, build_directory, verbose, with_cuda, is_python_module, is_standalone, keep_intermediates)
   1089         is_python_module,
   1090         is_standalone,
-> 1091         keep_intermediates=keep_intermediates)
   1092 
   1093 

d:\conda\envs\torchdrug\lib\site-packages\torch\utils\cpp_extension.py in _jit_compile(name, sources, extra_cflags, extra_cuda_cflags, extra_ldflags, extra_include_paths, build_directory, verbose, with_cuda, is_python_module, is_standalone, keep_intermediates)
   1315         return _get_exec_path(name, build_directory)
   1316 
-> 1317     return _import_module_from_library(name, build_directory, is_python_module)
   1318 
   1319 

d:\conda\envs\torchdrug\lib\site-packages\torch\utils\cpp_extension.py in _import_module_from_library(module_name, path, is_python_module)
   1697 def _import_module_from_library(module_name, path, is_python_module):
   1698     # https://stackoverflow.com/questions/67631/how-to-import-a-module-given-the-full-path
-> 1699     file, path, description = imp.find_module(module_name, [path])
   1700     # Close the .so file after load.
   1701     with file:

d:\conda\envs\torchdrug\lib\imp.py in find_module(name, path)
    294         break  # Break out of outer loop when breaking out of inner loop.
    295     else:
--> 296         raise ImportError(_ERR_MSG.format(name), name=name)
    297 
    298     encoding = None

ImportError: No module named 'torch_ext'

The same code works well in Colab and I suspect this is because I couldn’t install rdkit-pypi and installed rdkit on conda instead.

[KiddoZhu]

Thanks for the observation. TorchDrug is mainly developed and tested on Linux, and should work with Colab. Currently there is no guarantee for installing TorchDrug on Windows, but we will come up with Windows support soon.

[chaitjo]

Had the same issue, and was able to fix it by:

rm -rf /home/your_user_name/.cache/torch_extensions