[Port Hil-SERL] Add unit tests for the reward classifier & fix imports & check script (#578)

2026-07-25 10:46:01 +00:00 · 2024-12-23 16:43:55 +07:00
parent 7b68bfb73b
commit 70b652f791
7 changed files with 499 additions and 2 deletions
@@ -13,7 +13,7 @@ class ClassifierConfig:
    hidden_dim: int = 256
    dropout_rate: float = 0.1
    model_name: str = "microsoft/resnet-50"
-    device: str = "cuda" if torch.cuda.is_available() else "mps"
+    device: str = "cpu"
    model_type: str = "cnn"  # "transformer" or "cnn"

    def save_pretrained(self, save_dir):