Import error when trying to run TPU accelerator on CPU #12044
Description
🐛 Bug
To Reproduce
import os
import torch
from torch.utils.data import DataLoader, Dataset
from pytorch_lightning import LightningModule, Trainer
class RandomDataset(Dataset):
def __init__(self, size, length):
self.len = length
self.data = torch.randn(length, size)
def __getitem__(self, index):
return self.data[index]
def __len__(self):
return self.len
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)
def run():
train_data = DataLoader(RandomDataset(32, 64), batch_size=2)
val_data = DataLoader(RandomDataset(32, 64), batch_size=2)
test_data = DataLoader(RandomDataset(32, 64), batch_size=2)
model = BoringModel()
trainer = Trainer(
default_root_dir=os.getcwd(),
limit_train_batches=1,
limit_val_batches=1,
limit_test_batches=1,
num_sanity_val_steps=0,
max_epochs=1,
enable_model_summary=False,
accelerator="tpu",
devices=1,
)
trainer.fit(model, train_dataloaders=train_data, val_dataloaders=val_data)
trainer.test(model, dataloaders=test_data)
if __name__ == "__main__":
run()Produces
Traceback (most recent call last):
File "/Users/adrian/repositories/pytorch-lightning/pl_examples/bug_report/bug_report_model.py", line 68, in <module>
run()
File "/Users/adrian/repositories/pytorch-lightning/pl_examples/bug_report/bug_report_model.py", line 52, in run
trainer = Trainer(
File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/utilities/argparse.py", line 336, in insert_env_defaults
return fn(self, **kwargs)
File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/trainer/trainer.py", line 477, in __init__
self._accelerator_connector = AcceleratorConnector(
File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/trainer/connectors/accelerator_connector.py", line 194, in __init__
self._strategy_flag = self._choose_strategy()
File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/trainer/connectors/accelerator_connector.py", line 531, in _choose_strategy
return SingleTPUStrategy(device=self._parallel_devices[0]) # type: ignore
File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/strategies/single_tpu.py", line 44, in __init__
device=xm.xla_device(device),
NameError: name 'xm' is not defined
Expected behavior
Should error with
File "/Users/adrian/repositories/pytorch-lightning/pytorch_lightning/trainer/connectors/accelerator_connector.py", line 218, in select_accelerator_type
raise MisconfigurationException(f"You passed `accelerator='tpu'`, but {msg}.")
pytorch_lightning.utilities.exceptions.MisconfigurationException: You passed `accelerator='tpu'`, but TPUs are not available.
like on 1.5.10
Environment
Latest master.
Ask if you need more.