Spaces:

ReacherTN
/

ai_image_gen_detector

Runtime error

ReacherTN commited on Apr 21, 2023

Commit

d383e61

1 Parent(s): 8693834

Update predict.py

Files changed (1) hide show

predict.py CHANGED Viewed

@@ -10,7 +10,7 @@ def predict_one_image(path) :
     image = read_image(path)
     image = get_valid_augs()(image=image)['image']
     image = torch.tensor(image,dtype=torch.float)
-    image = image.reshape((1,3,224,224))
     model = CustomModel()
     #loading ckpt
     model.load_state_dict(torch.load(CKPT,map_location=torch.device('cpu')))

     image = read_image(path)
     image = get_valid_augs()(image=image)['image']
     image = torch.tensor(image,dtype=torch.float)
+    image = image.reshape((1,3,512,512))
     model = CustomModel()
     #loading ckpt
     model.load_state_dict(torch.load(CKPT,map_location=torch.device('cpu')))