Upload train_k.py

train_k.py

@@ -156,8 +156,10 @@ for j in range(1, 179+1):
         transform  = transforms
     )
 
-
-    model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained(config.ENCODER, config.DECODER)
+    if os.path.exists('VIT_large_gpt2_model'):
+        model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained('VIT_large_gpt2_model')
+    else:
+        model = VisionEncoderDecoderModel.from_encoder_decoder_pretrained(config.ENCODER, config.DECODER)
 
 
     model.config.decoder_start_token_id = tokenizer.cls_token_id