deepfake / deepfakeconfig.py
not-lain's picture
commit files to HF hub
2374ae7
raw
history blame
260 Bytes
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'