diff --git a/deckard/base/model/art_pipeline.py b/deckard/base/model/art_pipeline.py index 6b540f2d..aece4fb5 100644 --- a/deckard/base/model/art_pipeline.py +++ b/deckard/base/model/art_pipeline.py @@ -66,6 +66,12 @@ def __call__(self): import torch device_type = "gpu" if torch.cuda.is_available() else "cpu" + if device_type == "gpu": + logger.info("Using GPU") + logger.info("Model moved to GPU") + device = torch.device("cuda") + model.to(device) + data = [d.to(device) for d in data] model = TorchInitializer( data=data, model=model, diff --git a/deckard/base/model/model.py b/deckard/base/model/model.py index d335ef60..5e96557a 100644 --- a/deckard/base/model/model.py +++ b/deckard/base/model/model.py @@ -110,9 +110,6 @@ def __init__(self, **kwargs): logger.info(f"Initializing model trainer with kwargs {kwargs}") self.kwargs = kwargs - # def __hash__(self): - # return int(my_hash(self), 16) - def __call__(self, data: list, model: object, library=None): logger.info(f"Training model {model} with fit params: {self.kwargs}") device = str(model.device) if hasattr(model, "device") else "cpu" diff --git a/test/pipelines/evasion/.gitignore b/test/pipelines/evasion/.gitignore new file mode 100644 index 00000000..47f39c2f --- /dev/null +++ b/test/pipelines/evasion/.gitignore @@ -0,0 +1 @@ +.dvc/*