Fix model type detection in tests

Co-authored-by: Alexander Borzunov <[email protected]>

tests/test_optimized_layers.py

@@ -177,10 +177,6 @@ def _reorder_cache_from_llama_to_bloom(
         return (key_states, value_states)
 
 
-@pytest.mark.skipif(
-    all(model_name not in MODEL_NAME for model_name in ("falcon", "llama")),
-    reason="This test is applicable only to Falcon and LLaMa models",
-)
 @pytest.mark.parametrize("device", ["cpu", "cuda:0"])
 @pytest.mark.forked
 def test_falcon(device):
@@ -196,10 +192,12 @@ def test_falcon(device):
     block = config.block_class(config).to(dtype)
     block = convert_block(block, 0, config, tensor_parallel_devices, device, quant_type=quant_type, freeze=True)
 
-    if "falcon" in MODEL_NAME:
+    if config.model_type == "falcon":
         unopt_block = UnoptimizedWrappedFalconBlock(config).to(dtype)
-    elif "llama" in MODEL_NAME:
+    elif config.model_type == "llama":
         unopt_block = UnoptimizedWrappedLlamaBlock(config).to(dtype)
+    else:
+        pytest.skip(f"This test is not applicable to {config.model_type} models")
 
     unopt_block = convert_block(
         unopt_block, 0, config, tensor_parallel_devices, device, quant_type=quant_type, freeze=True