fixed code v2
This commit is contained in:
@@ -1,4 +1,5 @@
|
||||
import torch
|
||||
import os
|
||||
import torch
|
||||
from transformers import AutoImageProcessor
|
||||
from transformers import AutoModelForImageClassification
|
||||
from transformers import AutoModelForSequenceClassification
|
||||
@@ -17,6 +18,11 @@ class ModelManager:
|
||||
self.image_model = None
|
||||
self.image_processor = None
|
||||
|
||||
def _setup_hf_token(self):
|
||||
if settings.HF_TOKEN:
|
||||
os.environ["HF_TOKEN"] = settings.HF_TOKEN
|
||||
os.environ["HUGGING_FACE_HUB_TOKEN"] = settings.HF_TOKEN
|
||||
|
||||
def load_device(self):
|
||||
if settings.DEVICE == "auto":
|
||||
self.device = torch.device(
|
||||
|
||||
Reference in New Issue
Block a user