match config class

generation_demo.py

@@ -7,7 +7,7 @@ os.environ['TOKENIZERS_PARALLELISM'] = 'false'
 tokenizer = AutoTokenizer.from_pretrained("Caiyun-AI/DCFormer-2.8B")
 model = AutoModelForCausalLM.from_pretrained("Caiyun-AI/DCFormer-2.8B", trust_remote_code=True)
 
-device = torch.device('cuda')
+device = torch.device('cuda:1')
 MAX_BATCH_SIZE = 1
 MAX_SEQ_LENGTH = 2048
 NUM_TOKENS_TO_GENERATE = 100

modeling_dcformer.py

@@ -70,6 +70,7 @@ class KVKWCache(nn.Module):
         return k_out, v_out, kw_out 
 
 class DCFormer(PreTrainedModel):
+    config_class=DCFormerConfig
     '''
     DCFormer's implementation is adapted from https://github.com/pytorch-labs/gpt-fast/blob/main/model.py#L89 
     '''