title: Hugging Face model gives 0 accuracy after loading
body: <p>I'm trying to train a model with google's <a href="https://huggingface.co/google/vit-base-patch16-224" rel="nofollow noreferrer">ViT model</a> and an extra layer on a doodle dataset. After 5 hours of training, test accuracy increased from 0.0 to 0.75. However, when I saved and later loaded the model, test accuracy had fallen back to 0. This has happened 2-3 times. Could someone help me figure out what I am doing wrong?</p>
  <p>I was using Google's Quick Draw dataset downloaded from Google's <a href="https://www.kaggle.com/c/quickdraw-doodle-recognition" rel="nofollow noreferrer">Kaggle Competition</a></p>
  <p>Please find the necessary parts of my code here:</p>
  <pre><code>from torch import nn
  from pytorch_lightning import Trainer
  from pytorch_lightning import LightningModule
  from transformers import ViTModel, ViTConfig
  from transformers import ViTForImageClassification, AdamW
  from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
  
  model_name = 'google/vit-base-patch16-224-in21k'
  
  #################################################################
  #################### Model Class ################################
  #################################################################
  
  class DoodleTransformer(LightningModule):
      def __init__(self, num_labels):
          super(DoodleTransformer, self).__init__()
          
          self.vit = ViTModel.from_pretrained(model_name)
          self.classifier = nn.Linear(self.vit.config.hidden_size, num_labels)
  
      def forward(self, input_data):
          outputs = self.vit(pixel_values=input_data)
          logits = self.classifier(outputs.last_hidden_state[:, 0])
          
          return logits
      
      def common_step(self, batch, batch_idx):
          pixel_values = batch['input_data']
          labels = batch['labels']
          logits = self(pixel_values)
  
          criterion = nn.CrossEntropyLoss()
          loss = criterion(logits, labels)
          predictions = logits.argmax(-1)
          correct = (predictions == labels).sum().item()
          accuracy = correct/pixel_values.shape[0]
  
          return loss, accuracy
        
      def training_step(self, batch, batch_idx):
          pixel_values = batch['input_data']
          labels = batch['labels']
          logits = self(pixel_values)
  
          criterion = nn.CrossEntropyLoss()
          loss = criterion(logits, labels)
          predictions = logits.argmax(-1)
          correct = (predictions == labels).sum().item()
          accuracy = correct/pixel_values.shape[0]
          # logs metrics for each training_step,
          # and the average across the epoch
          self.log(&quot;accuracy&quot;, accuracy, on_step=True, on_epoch=False, prog_bar=True, logger=True)
  
          return loss
      
      def validation_step(self, batch, batch_idx):
          loss, accuracy = self.common_step(batch, batch_idx)
          
          self.log(&quot;validation_loss&quot;, loss)
          self.log(&quot;validation_accuracy&quot;, accuracy)
  
          return loss
  
      def test_step(self, batch, batch_idx):
          loss, accuracy = self.common_step(batch, batch_idx)
          
          self.log(&quot;test_loss&quot;, loss)
          self.log(&quot;test_accuracy&quot;, accuracy)
  
          return loss
  
      def configure_optimizers(self):
          # We could make the optimizer more fancy by adding a scheduler and specifying which parameters do
          # not require weight_decay but just using AdamW out-of-the-box works fine
          return AdamW(self.parameters(), lr=5e-5)
  
      def train_dataloader(self):
          return train_dataloader
  
      def val_dataloader(self):
          return val_dataloader
  
      def test_dataloader(self):
          return test_dataloader
      
   
  model_path = &quot;path/to/model.ckpt
  model_path_pt = &quot;path/to/state_dict.pt&quot;
  
  # I tried with both lines below as well as each one separately.
  model = DoodleTransformer.load_from_checkpoint(model_path, num_labels=340)
  model.load_state_dict(torch.load(model_path_pt))
  
  
  
  #################################################################
  ######################## Trainer ################################
  #################################################################
  
  early_stop_callback = EarlyStopping(
      monitor='val_loss',
      patience=1,
      strict=False,
      verbose=False,
      mode='min'
  )
  
  checkpoint_callback = ModelCheckpoint(
      monitor=&quot;validation_loss&quot;,
      dirpath=&quot;./model&quot;,
      filename=&quot;qd_model-3-{val_loss:.2f}&quot;,
      every_n_epochs=1,
      save_weights_only=True
  )
  
  trainer = Trainer(
      gpus=1,
      callbacks=[early_stop_callback, checkpoint_callback],
      max_epochs=1
  )
  
  #################################################################
  ######################## Training ###############################
  #################################################################
  
  # Work as expected
  trainer.fit(model)
  
  #################################################################
  ######################## Testing ################################
  #################################################################
  
  # Accuracy of ~0.75 after training. But -0.00 if loaded from saved or new model.
  test_results = trainer.test(model=model, dataloaders = test_dataloader)
  </code></pre>