From dc3b5ad82ea7c6c4651c4e2183414a6a19dd91f2 Mon Sep 17 00:00:00 2001
From: Zach Mueller <muellerzr@gmail.com>
Date: Fri, 9 Aug 2024 15:35:25 -0400
Subject: [PATCH] Fix deepspeed tests (#3003)

* Unpin deepspeed

* Include proper branch for docker image

* Properly working

* Revert all other changes
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index fb7d133f111..fdbc37d3f2d 100644
--- a/setup.py
+++ b/setup.py
@@ -36,7 +36,7 @@
     "timm",
 ]
 extras["testing"] = extras["test_prod"] + extras["test_dev"]
-extras["deepspeed"] = ["deepspeed<=0.14.0"]
+extras["deepspeed"] = ["deepspeed"]
 extras["rich"] = ["rich"]
 
 extras["test_trackers"] = ["wandb", "comet-ml", "tensorboard", "dvclive"]