From 9e0aeb5429a675291568f8ddc20d5dd8ba5b8afd Mon Sep 17 00:00:00 2001
From: EvenStrangest <github@mentor.org.il>
Date: Sat, 28 Sep 2024 05:55:43 +0300
Subject: [PATCH] Allow pytorch 2.0 & pytorch-lightning 2.0 via Copilot
 Workspace

Fixes #54

Update the package to support PyTorch 2.0 and PyTorch-Lightning 2.0.

* **pyproject.toml**
  - Update `torch` version to "^2.0.0".
  - Update `pytorch-lightning` version to "^2.0.0".

* **.github/workflows/ci.yml**
  - Add Python 3.11 to the matrix of Python versions for unit tests.

* **pycave/bayes/gmm/lightning_module.py**
  - Update import for `EarlyStopping` to ensure compatibility with PyTorch-Lightning 2.0.

* **pycave/bayes/markov_chain/estimator.py**
  - Add import for `Trainer` from `pytorch_lightning`.

* **pycave/utils/lightning_module.py**
  - Update import for `pytorch_lightning` to `lightning.pytorch` to ensure compatibility with PyTorch-Lightning 2.0.

---

For more details, open the [Copilot Workspace session](https://copilot-workspace.githubnext.com/borchero/pycave/issues/54?shareId=XXXX-XXXX-XXXX-XXXX).
---
 .github/workflows/ci.yml               | 2 +-
 pycave/bayes/gmm/lightning_module.py   | 2 +-
 pycave/bayes/markov_chain/estimator.py | 1 +
 pycave/utils/lightning_module.py       | 2 +-
 pyproject.toml                         | 4 ++--
 5 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index aa339e7..efea437 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -27,7 +27,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        python-version: ["3.8", "3.9", "3.10"]
+        python-version: ["3.8", "3.9", "3.10", "3.11"]
     steps:
       - name: Checkout
         uses: actions/checkout@v3
diff --git a/pycave/bayes/gmm/lightning_module.py b/pycave/bayes/gmm/lightning_module.py
index a5aaf90..27b33a0 100644
--- a/pycave/bayes/gmm/lightning_module.py
+++ b/pycave/bayes/gmm/lightning_module.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 import pytorch_lightning as pl
 import torch
-from pytorch_lightning.callbacks import EarlyStopping
+from pytorch_lightning.callbacks.early_stopping import EarlyStopping
 from torchmetrics import MeanMetric
 from pycave.bayes.core import cholesky_precision
 from pycave.utils import NonparametricLightningModule
diff --git a/pycave/bayes/markov_chain/estimator.py b/pycave/bayes/markov_chain/estimator.py
index 11fc167..12ef1d7 100644
--- a/pycave/bayes/markov_chain/estimator.py
+++ b/pycave/bayes/markov_chain/estimator.py
@@ -7,6 +7,7 @@
 from lightkit.data import DataLoader, dataset_from_tensors
 from torch.nn.utils.rnn import PackedSequence
 from torch.utils.data import Dataset
+from pytorch_lightning import Trainer
 from .lightning_module import MarkovChainLightningModule
 from .model import MarkovChainModel, MarkovChainModelConfig
 from .types import collate_sequences, collate_sequences_same_length, SequenceData
diff --git a/pycave/utils/lightning_module.py b/pycave/utils/lightning_module.py
index 049d4fa..15cda0d 100644
--- a/pycave/utils/lightning_module.py
+++ b/pycave/utils/lightning_module.py
@@ -1,6 +1,6 @@
 from abc import ABC, abstractmethod
 from typing import List
-import pytorch_lightning as pl
+import lightning.pytorch as pl
 import torch
 from torch import nn
 
diff --git a/pyproject.toml b/pyproject.toml
index 17146fa..7ab0868 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -16,8 +16,8 @@ version = "0.0.0"
 lightkit = "^0.5.0"
 numpy = "^1.20.3"
 python = ">=3.8,<3.11"
-pytorch-lightning = "^1.6.0"
-torch = "^1.8.0"
+pytorch-lightning = "^2.0.0"
+torch = "^2.0.0"
 torchmetrics = ">=0.6,<0.12"
 
 [tool.poetry.group.pre-commit.dependencies]