fork-detect-fake / deepfakeconfig.py
wyyadd's picture
feat: update pipeline
d30298c
from transformers import PretrainedConfig
import torch
class DeepFakeConfig(PretrainedConfig):
model_type = "ResNet"
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.DEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'