Update InpaintReward.py

InpaintReward.py

@@ -54,7 +54,7 @@ class ViTBlock(nn.Module):
         return x
 
 
-class ImageReward(nn.Module):
+class InpaintReward(nn.Module):
     def __init__(self, config, device='cpu'):
         super().__init__()
         self.config = config
@@ -62,7 +62,7 @@ class ImageReward(nn.Module):
         
         self.clip_model, self.preprocess = clip.load("ViT-B/32") 
         self.clip_model = self.clip_model.float()
-        self.mlp = MLP(self.config['ImageReward']['mlp_dim'])
+        self.mlp = MLP(self.config['Reward']['mlp_dim'])
         self.vit_block = ViTBlock(self.config["ViT"]["feature_dim"], self.config["ViT"]["num_heads"], self.config["ViT"]["mlp_dim"])
         
         self.toImage = transforms.ToPILImage()
@@ -127,7 +127,7 @@ class ImageReward(nn.Module):
 
        
 
-class ImageRewardGroup(nn.Module):
+class InpaintRewardGroup(nn.Module):
     def __init__(self, config, device='cpu'):
         super().__init__()
         self.config = config
@@ -136,7 +136,7 @@ class ImageRewardGroup(nn.Module):
         self.clip_model, self.preprocess = clip.load("ViT-B/32", device="cuda") #clip.load(config['clip_model'], device="cuda" if torch.cuda.is_available() else "cpu")
         
         self.clip_model = self.clip_model.float()
-        self.mlp = MLP(config['ImageReward']['mlp_dim'])
+        self.mlp = MLP(config['Reward']['mlp_dim'])
         self.vit_block = ViTBlock(self.config["ViT"]["feature_dim"], self.config["ViT"]["num_heads"], self.config["ViT"]["mlp_dim"])
         
         if self.config.fix_base: