(self)
| 59 | @require_torch_gpu |
| 60 | @require_flashpack |
| 61 | def test_load_model_device(self): |
| 62 | model = AutoModel.from_pretrained(self.model_id, subfolder="transformer") |
| 63 | with tempfile.TemporaryDirectory() as temp_dir: |
| 64 | model.save_pretrained(temp_dir, use_flashpack=True) |
| 65 | model = AutoModel.from_pretrained(temp_dir, use_flashpack=True, device_map={"": torch.device("cuda")}) |
| 66 | self.assertTrue(model.device.type == "cuda") |
| 67 | |
| 68 | @require_flashpack |
| 69 | def test_load_model_device_auto(self): |
nothing calls this directly
no test coverage detected