code: def test_prepare(self):
      pre_train_config = self._get_pre_train_config()
      pre_train_task = build_task(pre_train_config)
      pre_train_task.prepare()
      checkpoint = get_checkpoint_dict(pre_train_task, {})
  
      fine_tuning_config = self._get_fine_tuning_config()
      fine_tuning_task = build_task(fine_tuning_config)
  
      # test: cannot prepare a fine tuning task without a pre-trained checkpoint
      with self.assertRaises(Exception):
          fine_tuning_task.prepare()
  
      # test: prepare should succeed after pre-trained checkpoint is set
      fine_tuning_task = build_task(fine_tuning_config)
      fine_tuning_task._set_pretrained_checkpoint_dict(checkpoint)
      fine_tuning_task.prepare()
  
      # test: prepare should succeed if a pre-trained checkpoint is provided in the
      # config
      fine_tuning_config = self._get_fine_tuning_config(
          pretrained_checkpoint=True
      )
      fine_tuning_task = build_task(fine_tuning_config)
      with mock.patch(
          "classy_vision.tasks.fine_tuning_task.load_and_broadcast_checkpoint",
          return_value=checkpoint,
      ):
          fine_tuning_task.prepare()
  
      # test: a fine tuning task with incompatible heads with a manually set
      # pre-trained checkpoint should fail to prepare if the heads are not reset
      fine_tuning_config = self._get_fine_tuning_config(
          head_num_classes=10
      )
      fine_tuning_task = build_task(fine_tuning_config)
      fine_tuning_task._set_pretrained_checkpoint_dict(checkpoint)
  
      with self.assertRaises(Exception):
          fine_tuning_task.prepare()
  
      # test: a fine tuning task with incompatible heads with a manually set
      # pre-trained checkpoint should succeed to prepare if the heads are reset
      fine_tuning_task = build_task(fine_tuning_config)
      fine_tuning_task._set_pretrained_checkpoint_dict(
          copy.deepcopy(checkpoint)
      ).set_reset_heads(True)
  
      fine_tuning_task.prepare()
  
      # test: a fine tuning task with incompatible heads with the pre-trained
      # checkpoint provided in the config should fail to prepare
      fine_tuning_config = self._get_fine_tuning_config(
          head_num_classes=10, pretrained_checkpoint=True
      )
  
      fine_tuning_task = build_task(fine_tuning_config)
      with mock.patch(
          "classy_vision.tasks.fine_tuning_task.load_and_broadcast_checkpoint",
          return_value=copy.deepcopy(checkpoint),
      ) and self.assertRaises(Exception):
          fine_tuning_task.prepare()
  
      # test: a fine tuning task with incompatible heads with the pre-trained
      # checkpoint provided in the config should succeed to prepare if the heads are
      # reset
      fine_tuning_task = build_task(fine_tuning_config)
      fine_tuning_task.set_reset_heads(True)
      with mock.patch(
          "classy_vision.tasks.fine_tuning_task.load_and_broadcast_checkpoint",
          return_value=copy.deepcopy(checkpoint),
      ):
          fine_tuning_task.prepare()
  
imports: import copy
  import unittest
  from typing import Any, Dict
  from unittest import mock
  import torch
  import torch.nn
  from classy_vision.generic.util import get_checkpoint_dict
  from classy_vision.losses import ClassyLoss, register_loss
  from classy_vision.tasks import build_task, FineTuningTask
  from classy_vision.trainer import LocalTrainer
  from test.generic.config_utils import get_fast_test_task_config
  from test.generic.utils import compare_model_state
  from torch.nn.modules.loss import CrossEntropyLoss
label_desc: Use image node_id 6 for calling the TestFineTuningTask obj's underlying member method code with example usage: obj.test_prepare() without return types