From 057c4ce67f32f89b7b6bcdb9e34fdc09fbb68a03 Mon Sep 17 00:00:00 2001
From: Alexander Suslov <alexander.suslov@intel.com>
Date: Wed, 24 Apr 2024 16:33:58 +0400
Subject: [PATCH] Checking model type in nncf.torch.wrap_model (#2651)

### Changes

Added model type check in nncf.torch.wrap_model

### Reason for changes

Improve UX

### Related tickets

https://github.com/openvinotoolkit/nncf/pull/2636#discussion_r1576234452

### Tests

N/A
---
 nncf/torch/model_creation.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/nncf/torch/model_creation.py b/nncf/torch/model_creation.py
index ae5386aa7e7..b6d21cc66a9 100644
--- a/nncf/torch/model_creation.py
+++ b/nncf/torch/model_creation.py
@@ -326,6 +326,11 @@ def wrap_model(
     :param trace_parameters: Whether to trace model parameters. Default is False.
     :return: A model wrapped by NNCFNetwork.
     """
+    if not isinstance(model, torch.nn.Module):
+        raise TypeError(
+            f"The provided model type {type(model)} is incompatible. "
+            "Only models inheriting from torch.nn.Module are supported."
+        )
 
     input_info = ExampleInputInfo.from_example_input(example_input)