Upload extract.py

extract.py

@@ -7,11 +7,10 @@ import torch
 from torchgeo.models import dofa_base_patch16_224
 
 # Load the checkpoint
-in_filename = "ofa_base_checkpoint_e99.pth"
-checkpoint = torch.load(in_filename, map_location=torch.device("cpu"))
+in_filename = "DOFA_ViT_base_e100.pth"
+weights = torch.load(in_filename, map_location=torch.device("cpu"))
 
 # Remove extra keys
-weights = checkpoint["model"]
 del weights["mask_token"]
 del weights["norm.weight"], weights["norm.bias"]
 del weights["projector.weight"], weights["projector.bias"]