def create_and_check_model(self, config, input_ids, token_type_ids,
                            sequence_labels, token_labels, choice_labels):
     model = FNetModel(config=config)
     model.to(torch_device)
     model.eval()
     result = model(input_ids, token_type_ids=token_type_ids)
     result = model(input_ids)
     self.parent.assertEqual(
         result.last_hidden_state.shape,
         (self.batch_size, self.seq_length, self.hidden_size))
    def test_inference_model(self):
        model = FNetModel.from_pretrained("google/fnet-base")
        model.to(torch_device)

        input_ids = torch.tensor([[0, 1, 2, 3, 4, 5]], device=torch_device)
        output = model(input_ids)[0]

        expected_shape = torch.Size((1, 6, model.config.hidden_size))
        self.assertEqual(output.shape, expected_shape)

        expected_slice = torch.tensor(
            [[[4.1541, -0.1051, -0.1667], [-0.9144, 0.2939, -0.0086],
              [-0.8472, -0.7281, 0.0256]]],
            device=torch_device)

        self.assertTrue(
            torch.allclose(output[:, :3, :3], expected_slice, atol=1e-4))
 def test_model_from_pretrained(self):
     for model_name in FNET_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
         model = FNetModel.from_pretrained(model_name)
         self.assertIsNotNone(model)