Fix ModuleNotFoundError: entry scripts import lightning.pytorch but environment.yml ships pytorch-lightning

[alex-crlhmmr]

Summary

pretrain.py, classification.py, and segmentation.py import seed_everything from lightning.pytorch, but environment.yml only installs pytorch-lightning (a different PyPI distribution, not lightning). As a result, the documented setup path fails at import time: after conda env create -f environment.yml, any training run raises ModuleNotFoundError: No module named 'lightning'.

This unifies the three imports onto pytorch_lightning, which the rest of the codebase and environment.yml already use.

Steps to reproduce

conda env create -f environment.yml
conda activate brep2shape
python pretrain.py train --dataset_dir <...>

Traceback (most recent call last):
  File "pretrain.py", line 102, in <module>
    from lightning.pytorch import seed_everything
ModuleNotFoundError: No module named 'lightning'

Confirmed in a clean environment: pip install pytorch-lightning==2.2.4 (exactly what environment.yml specifies) does not install the lightning package, and lightning is not among its dependencies (Requires: fsspec, lightning-utilities, numpy, packaging, PyYAML, torch, torchmetrics, tqdm, typing-extensions).

Root cause

The codebase mixes the two Lightning distributions. The core already uses pytorch_lightning:

• utils/training.py — Trainer, ModelCheckpoint, TensorBoardLogger, DDPStrategy/FSDPStrategy
• models/*.py — import pytorch_lightning as pl

Only three call sites reach for the other distribution:

• pretrain.py:102
• classification.py:99
• segmentation.py:114

each from lightning.pytorch import seed_everything. Since environment.yml pins pytorch-lightning==2.2.4 and never installs lightning, these three lines break the shipped environment. pytorch_lightning re-exports seed_everything as a top-level symbol (verified in 2.2.4), so unifying on it is behavior-preserving and avoids requiring two overlapping Lightning distributions.

Fix

- from lightning.pytorch import seed_everything
+ from pytorch_lightning import seed_everything

in pretrain.py, classification.py, and segmentation.py.

Testing

• from pytorch_lightning import seed_everything resolves in the unmodified environment.yml environment (pytorch-lightning 2.2.4, Python 3.10.8, PyTorch 2.2.1, CUDA 12.1, DGL 2.1.0).
• seed_everything is the only symbol imported from lightning.pytorch anywhere in the repo, so no other call sites are affected.