Update app.py

app.py

@@ -164,7 +164,7 @@ valid_dataset = CustomOCRDataset(
 )
 
 model = VisionEncoderDecoderModel.from_pretrained(ModelConfig.MODEL_NAME)
-model.to(device)
+#model.to(device)
 print(model)
 # Total parameters and trainable parameters.
 total_params = sum(p.numel() for p in model.parameters())
@@ -214,7 +214,7 @@ training_args = Seq2SeqTrainingArguments(
     evaluation_strategy='epoch',
     per_device_train_batch_size=TrainingConfig.BATCH_SIZE,
     per_device_eval_batch_size=TrainingConfig.BATCH_SIZE,
-    fp16=True,
+    #fp16=True,
     output_dir='seq2seq_model_printed/',
     logging_strategy='epoch',
     save_strategy='epoch',
@@ -237,7 +237,8 @@ trainer = Seq2SeqTrainer(
 res = trainer.train()
 
 processor = TrOCRProcessor.from_pretrained(ModelConfig.MODEL_NAME)
-trained_model = VisionEncoderDecoderModel.from_pretrained('seq2seq_model_printed/checkpoint-'+str(res.global_step)).to(device)
+trained_model = VisionEncoderDecoderModel.from_pretrained('seq2seq_model_printed/checkpoint-'+str(res.global_step))
+#.to(device)
 
 def read_and_show(image_path):
     """
@@ -261,7 +262,8 @@ def ocr(image, processor, model):
         generated_text: the OCR'd text string.
     """
     # We can directly perform OCR on cropped images.
-    pixel_values = processor(image, return_tensors='pt').pixel_values.to(device)
+    pixel_values = processor(image, return_tensors='pt').pixel_values
+    #.to(device)
     generated_ids = model.generate(pixel_values)
     generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return generated_text