ImportError after installing torchtext 0.11.0 with conda

I have installed pytorch version 1.10.0 alongside torchtext, torchvision and torchaudio using conda. My PyTorch is cpu-only, and I have experimented with both conda install pytorch-mutex -c pytorch and conda install pytorch cpuonly -c pytorch to install the cpuonly version, both yielding the same eror that I will describe in the following lines.

I have also installed pytorch-lightning in conda, alongside jsonargparse[summaries via pip in the environment.

I have written this code to see whether LightningCLI works or not.

# script.py
import torch 
import pytorch_lightning as pl

class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def training_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

    def validation_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("valid_loss", loss)

    def test_step(self, batch, batch_idx):
        loss = self(batch).sum()
        self.log("test_loss", loss)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)

cli = LightningCLI(BoringModel)

But when I run it using python -m script fit --print_config, I get the following error:

ImportError: /home/farhood/miniconda3/envs/pytorch_dummy_environment/lib/python3.9/site-packages/torchtext/_torchtext.so: undefined symbol: _ZNK5torch3jit6MethodclESt6vectorIN3c106IValueESaIS4_EERKSt13unordered_mapISsS4_St4hashISsESt8equal_toISsESaISt4pairIKSsS4_EEE

Which indicates that there is something broken with my Conda installation, and it's probably related to torchtext somehow.

This is the versions of the installed torch related packages:

pytorch                   1.10.0          cpu_py39hc5866cc_0    conda-forge
pytorch-lightning         1.5.2              pyhd8ed1ab_0    conda-forge
pytorch-mutex             1.0                        cuda    pytorch
torchaudio                0.10.0                 py39_cpu    pytorch
torchmetrics              0.6.0              pyhd8ed1ab_0    conda-forge
torchtext                 0.11.0                     py39    pytorch
torchvision               0.11.1                 py39_cpu    pytorch

Solution

So in order to fix the problem, I had to change my environment.yaml in order to force pytorch to install from the pytorch channel.

So this is my environment.yaml now:

channels:
  - defaults
  - pytorch
  - conda-forge
dependencies:
  # ML section
  - pytorch::pytorch
  - pytorch::torchtext
  - pytorch::torchvision
  - pytorch::torchaudio
  - pytorch::cpuonly
  - mlflow=1.21.0
  - pytorch-lightning>=1.5.2
  - pip:
    - jsonargparse[signatures]

Using this I don't get the error anymore. The pytorch related stuff installed now is:

cpuonly                   2.0                           0    pytorch
pytorch                   1.10.0              py3.9_cpu_0    pytorch
pytorch-lightning         1.5.2              pyhd8ed1ab_0    conda-forge
pytorch-mutex             1.0                         cpu    pytorch
torchaudio                0.10.0                 py39_cpu  [cpuonly]  pytorch
torchtext                 0.11.0                     py39    pytorch
torchvision               0.11.1                 py39_cpu  [cpuonly]  pytorch