Update handler.py

handler.py

@@ -1,16 +1,17 @@
-from contextlib import nullcontext
 import time
 import torch
-from apogee.tokenizer import Tokenizer
-from apogee.model import GPT, ModelConfig
 
+from contextlib import nullcontext
 from typing import Any, Dict, Optional, Union
 from pathlib import Path
 
+from apogee.tokenizer import Tokenizer
+from apogee.model import GPT, ModelConfig
+
 torch.backends.cuda.matmul.allow_tf32 = True # allow tf32 on matmul
 torch.backends.cudnn.allow_tf32 = True # allow tf32 on cudnn
 
-class ApogeeHandler:
+class EndpointHandler:
     """
     Handler class.
     """
@@ -39,7 +40,7 @@ class ApogeeHandler:
         self.model.eval()
         self.model.to(self.device)
         self.model = torch.compile(self.model)
-        dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else 'float16' # 'float32' or 'bfloat16' or 'float16'
+        dtype = 'bfloat16' if torch.cuda.is_available() and torch.cuda.is_bf16_supported() and torch.cuda.get_device_capability()[0] >= 8 else 'float16' # 'float32' or 'bfloat16' or 'float16'
         ptdtype = {'float32': torch.float32, 'bfloat16': torch.bfloat16, 'float16': torch.float16}[dtype]
         self.ctx = nullcontext() if device == 'cpu' else torch.amp.autocast(device_type=device, dtype=ptdtype)
         print("Warming up hardware 🔥")
@@ -124,7 +125,7 @@ class ApogeeHandler:
 
 if __name__ == "__main__":
     import pandas as pd
-    handler = ApogeeHandler()
+    handler = EndpointHandler()
     test_path = Path(__file__).parents[2] / "tests" / "assets" / "BTCUSDT-1m-2019-03.csv"
     with open(test_path, "r") as f:
         data = pd.read_csv(f)