pinecone
/

ConstBERT

Feature Extraction

text-embeddings-inference

Model card Files Files and versions

caesar-one commited on May 14, 2025

Commit

624a902

·

verified ·

1 Parent(s): 39157e5

Upload ConstBERT

Files changed (1) hide show

modeling.py +2 -2

modeling.py CHANGED Viewed

@@ -21,10 +21,10 @@ class MixedPrecisionManager():
         self.activated = activated
         if self.activated:
-            self.scaler = torch.cuda.amp.GradScaler()
     def context(self):
-        return torch.cuda.amp.autocast() if self.activated else NullContextManager()
     def backward(self, loss):
         if self.activated:

         self.activated = activated
         if self.activated:
+            self.scaler = torch.amp.GradScaler("cuda")
     def context(self):
+        return torch.amp.autocast("cuda") if self.activated else NullContextManager()
     def backward(self, loss):
         if self.activated: