# Buggy code
def load(
    cls,
    meta: Dict[Text, Any],
    model_dir: Text = None,
    model_metadata: "Metadata" = None,
    cached_component: Optional["EmbeddingIntentClassifier"] = None,
    **kwargs: Any,
) -> "EmbeddingIntentClassifier":
    """Loads the trained model from the provided directory."""

    if model_dir and meta.get("file"):
        file_name = meta.get("file")
        checkpoint = os.path.join(model_dir, file_name + ".ckpt")

        with open(os.path.join(model_dir, file_name + ".tf_config.pkl"), "rb") as f:
            _tf_config = pickle.load(f)

        graph = tf.Graph()
        with graph.as_default():
            session = tf.compat.v1.Session(config=_tf_config)
            saver = tf.compat.v1.train.import_meta_graph(checkpoint + ".meta")

            saver.restore(session, checkpoint)

            batch_in = train_utils.load_tensor("batch_placeholder")

            sim_all = train_utils.load_tensor("similarity_all")
            pred_confidence = train_utils.load_tensor("pred_confidence")
            sim = train_utils.load_tensor("similarity")

            message_embed = train_utils.load_tensor("message_embed")
            label_embed = train_utils.load_tensor("label_embed")
            all_labels_embed = train_utils.load_tensor("all_labels_embed")

        with open(
            os.path.join(model_dir, file_name + ".inv_label_dict.pkl"), "rb"
        ) as f:
            inv_label_dict = pickle.load(f)

        with open(
            os.path.join(model_dir, file_name + ".batch_tuple_sizes.pkl"), "rb"
        ) as f:
            batch_tuple_sizes = pickle.load(f)

        return cls(
            component_config=meta,
            inverted_label_dict=inv_label_dict,
            session=session,
            graph=graph,
            batch_placeholder=batch_in,
            similarity_all=sim_all,
            pred_confidence=pred_confidence,
            similarity=sim,
            message_embed=message_embed,
            label_embed=label_embed,
            all_labels_embed=all_labels_embed,
            batch_tuple_sizes=batch_tuple_sizes,
        )

    else:
        raise_warning(
            f"Failed to load nlu model. "
            f"Maybe the path '{os.path.abspath(model_dir)}' doesn't exist?",
        )
        return cls(component_config=meta)


# Traceback (tensorflow.python.framework.errors_impl.InvalidArgumentError)
2020-01-30 16:57:43 INFO     rasa.nlu.utils.spacy_utils  - Trying to load spacy model with name 'de_core_news_sm'
2020-01-30 16:57:44 INFO     transformers.file_utils  - PyTorch version 1.3.0+cpu available.
2020-01-30 16:57:44 INFO     transformers.modeling_xlnet  - Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex .
2020-01-30 16:57:45 INFO     rasa.nlu.components  - Added 'SpacyNLP' to component cache. Key 'SpacyNLP-de_core_news_sm'.
2020-01-30 16:57:45 INFO     rasa.nlu.utils.spacy_utils  - Trying to load spacy model with name 'de_core_news_md'
2020-01-30 16:58:26 INFO     rasa.nlu.components  - Added 'SpacyNLP' to component cache. Key 'SpacyNLP-de_core_news_md'.
2020-01-30 16:58:26 INFO     rasa.nlu.training_data.training_data  - Training data stats:
- intent examples: 347 (18 distinct intents)
- Found intents: 'repeat_last_question', 'functions', 'deny+utter_mobilenumber', 'deny', 'utter_mobilenumber', 'confirm', 'reset_dialog', 'ask_outofscope', 'deny_appointment_suggestion', 'ask_repeat_contact', 'make_appointment', 'not_now', 'greet', 'goodbye', 'thankyou', 'ask_persisted_user_data', 'persist_purpose', 'select_appointment_slot'
- Number of response examples: 0 (0 distinct response)
- entity examples: 104 (7 distinct entities)
- found entities: 'appointment_type', 'daytime', 'not_now_reason', 'slot_shift', 'mobilenumber', 'time', 'employee'

2020-01-30 16:58:26 INFO     rasa.nlu.model  - Starting to train component SpacyNLP
2020-01-30 16:58:26 INFO     rasa.nlu.model  - Finished training component.
2020-01-30 16:58:26 INFO     rasa.nlu.model  - Starting to train component SpacyTokenizer
2020-01-30 16:58:26 INFO     rasa.nlu.model  - Finished training component.
2020-01-30 16:58:26 INFO     rasa.nlu.model  - Starting to train component CountVectorsFeaturizer
2020-01-30 16:58:26 INFO     rasa.nlu.model  - Finished training component.
2020-01-30 16:58:26 INFO     rasa.nlu.model  - Starting to train component CountVectorsFeaturizer
2020-01-30 16:58:27 INFO     rasa.nlu.model  - Finished training component.
2020-01-30 16:58:27 INFO     rasa.nlu.model  - Starting to train component SpacyFeaturizer
2020-01-30 16:58:27 INFO     rasa.nlu.model  - Finished training component.
2020-01-30 16:58:27 INFO     rasa.nlu.model  - Starting to train component EmbeddingIntentClassifier
Tensor("mul:0", shape=(?, ?, 192), dtype=float32)
Tensor("mul_1:0", shape=(?, ?, 20), dtype=float32)
Tensor("mul_2:0", shape=(18, 2, 20), dtype=float32)
Traceback (most recent call last):
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\framework\ops.py", line 1607, in _create_c_op
c_op = c_api.TF_FinishOperation(op_desc)
tensorflow.python.framework.errors_impl.InvalidArgumentError: Can not squeeze dim[1], expected a dimension of 1, got 2 for 'Squeeze_2' (op: 'Squeeze') with input shapes: [18,2,20].

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
File "c:\users\\appdata\local\programs\python\python36\Lib\runpy.py", line 193, in _run_module_as_main
"__main__", mod_spec)
File "c:\users\\appdata\local\programs\python\python36\Lib\runpy.py", line 85, in _run_code
exec(code, run_globals)
File "C:\Users\\Envs\eva_env\Scripts\rasa.exe\__main__.py", line 7, in <module>
File "c:\users\\envs\eva_env\lib\site-packages\rasa\__main__.py", line 76, in main
cmdline_arguments.func(cmdline_arguments)
File "c:\users\\envs\eva_env\lib\site-packages\rasa\cli\train.py", line 76, in train
additional_arguments=extract_additional_arguments(args),
File "c:\users\\envs\eva_env\lib\site-packages\rasa\train.py", line 50, in train
additional_arguments=additional_arguments,
File "c:\users\\appdata\local\programs\python\python36\Lib\asyncio\base_events.py", line 484, in run_until_complete
return future.result()
File "c:\users\\envs\eva_env\lib\site-packages\rasa\train.py", line 101, in train_async
additional_arguments,
File "c:\users\\envs\eva_env\lib\site-packages\rasa\train.py", line 188, in _train_async_internal
additional_arguments=additional_arguments,
File "c:\users\\envs\eva_env\lib\site-packages\rasa\train.py", line 245, in _do_training
persist_nlu_training_data=persist_nlu_training_data,
File "c:\users\\envs\eva_env\lib\site-packages\rasa\train.py", line 474, in _train_nlu_with_validated_data
persist_nlu_training_data=persist_nlu_training_data,
File "c:\users\\envs\eva_env\lib\site-packages\rasa\nlu\train.py", line 86, in train
interpreter = trainer.train(training_data, **kwargs)
File "c:\users\\envs\eva_env\lib\site-packages\rasa\nlu\model.py", line 191, in train
updates = component.train(working_data, self.config, **context)
File "c:\users\\envs\eva_env\lib\site-packages\rasa\nlu\classifiers\embedding_intent_classifier.py", line 769, in train
loss, acc = self._build_tf_train_graph(session_data)
File "c:\users\\envs\eva_env\lib\site-packages\rasa\nlu\classifiers\embedding_intent_classifier.py", line 571, in _build_tf_train_graph
label_data["label_features"], label_data["label_mask"][0], "label"
File "c:\users\\envs\eva_env\lib\site-packages\rasa\nlu\classifiers\embedding_intent_classifier.py", line 547, in _combine_sparse_dense_features
output = tf.squeeze(output, axis=1)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\util\dispatch.py", line 180, in wrapper
return target(*args, **kwargs)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\util\deprecation.py", line 507, in new_func
return func(*args, **kwargs)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\ops\array_ops.py", line 3651, in squeeze
return gen_array_ops.squeeze(input, axis, name)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\ops\gen_array_ops.py", line 10068, in squeeze
"Squeeze", input=input, squeeze_dims=axis, name=name)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\framework\op_def_library.py", line 794, in _apply_op_helper
op_def=op_def)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\util\deprecation.py", line 507, in new_func
return func(*args, **kwargs)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\framework\ops.py", line 3357, in create_op
attrs, op_def, compute_device)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\framework\ops.py", line 3426, in _create_op_internal
op_def=op_def)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\framework\ops.py", line 1770, in __init__
control_input_ops)
File "c:\users\\envs\eva_env\lib\site-packages\tensorflow_core\python\framework\ops.py", line 1610, in _create_c_op
raise ValueError(str(e))
ValueError: Can not squeeze dim[1], expected a dimension of 1, got 2 for 'Squeeze_2' (op: 'Squeeze') with input shapes: [18,2,20].

# Fixed code
def load(
    cls,
    meta: Dict[Text, Any],
    model_dir: Text = None,
    model_metadata: "Metadata" = None,
    cached_component: Optional["EmbeddingIntentClassifier"] = None,
    **kwargs: Any,
) -> "EmbeddingIntentClassifier":
    """Loads the trained model from the provided directory."""

    if model_dir and meta.get("file"):
        file_name = meta.get("file")
        checkpoint = os.path.join(model_dir, file_name + ".ckpt")

        with open(os.path.join(model_dir, file_name + ".tf_config.pkl"), "rb") as f:
            _tf_config = pickle.load(f)

        graph = tf.Graph()
        with graph.as_default():
            session = tf.compat.v1.Session(config=_tf_config)
            saver = tf.compat.v1.train.import_meta_graph(checkpoint + ".meta")

            saver.restore(session, checkpoint)

            batch_in = train_utils.load_tensor("batch_placeholder")

            sim_all = train_utils.load_tensor("similarity_all")
            pred_confidence = train_utils.load_tensor("pred_confidence")
            sim = train_utils.load_tensor("similarity")

            message_embed = train_utils.load_tensor("message_embed")
            label_embed = train_utils.load_tensor("label_embed")
            all_labels_embed = train_utils.load_tensor("all_labels_embed")

        with open(
            os.path.join(model_dir, file_name + ".inv_label_dict.pkl"), "rb"
        ) as f:
            inv_label_dict = pickle.load(f)

        with open(
            os.path.join(model_dir, file_name + ".batch_tuple_sizes.pkl"), "rb"
        ) as f:
            batch_tuple_sizes = pickle.load(f)

        return cls(
            component_config=meta,
            inverted_label_dict=inv_label_dict,
            session=session,
            graph=graph,
            batch_placeholder=batch_in,
            similarity_all=sim_all,
            pred_confidence=pred_confidence,
            similarity=sim,
            message_embed=message_embed,
            label_embed=label_embed,
            all_labels_embed=all_labels_embed,
            batch_tuple_sizes=batch_tuple_sizes,
        )

    else:
        raise_warning(
            f"Failed to load nlu model. "
            f"Maybe the path '{os.path.abspath(model_dir)}' doesn't exist?"
        )
        return cls(component_config=meta)