Gemma default device fix (#161)

* Add max_memory parameter to run config

Co-authored-by: Simon Schrader <[email protected]>

* Use configurable max_memory for offline explainer

Co-authored-by: Simon Schrader <[email protected]>

* Fix breaking change in prompt input formatting from vLLM

Co-authored-by: Simon Schrader <[email protected]>

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Fix bug in gemmascope device type check

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Simon Schrader <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: 3 people

delphi/sparse_coders/custom/gemmascope.py

@@ -104,6 +104,8 @@ def from_pretrained(cls, model_name_or_path, position, device):
         pt_params = {k: torch.from_numpy(v) for k, v in params.items()}
         model = cls(params["W_enc"].shape[0], params["W_enc"].shape[1])
         model.load_state_dict(pt_params)
-        if device == "cuda":
+        if device == "cuda" or (
+            isinstance(device, torch.device) and device.type == "cuda"
+        ):
             model.cuda()
         return model