Fix versions

mosaicml · Nov 22, 2024 · efb9c72 · efb9c72
1 parent 614f1a9
commit efb9c72
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/diffusion/models/transformer.py b/diffusion/models/transformer.py
@@ -192,7 +192,7 @@ def forward(self, x: torch.Tensor, t: torch.Tensor) -> Tuple[torch.Tensor, torch
         q, k, v = self.qkv(x).chunk(3, dim=-1)
         q = self.q_norm(q)
         k = self.k_norm(k)
-        return q.contiguous(), k.contiguous(), v.contiguous()
+        return q, k, v
 
 
 class SelfAttention(nn.Module):

diff --git a/setup.py b/setup.py
@@ -6,8 +6,8 @@
 from setuptools import find_packages, setup
 
 install_requires = [
-    'mosaicml>=0.27.0',
-    'mosaicml-streaming>=0.9.0',
+    'mosaicml==0.27.0',
+    'mosaicml-streaming==0.9.0',
     'hydra-core>=1.2',
     'hydra-colorlog>=1.1.0',
     'diffusers[torch]==0.30.3',