# Buggy code
def __init__(
    self,
    model_name: str,
    *,
    max_length: int = None,
    sub_module: str = None,
    train_parameters: bool = True,
    last_layer_only: bool = True,
    override_weights_file: Optional[str] = None,
    override_weights_strip_prefix: Optional[str] = None,
    gradient_checkpointing: Optional[bool] = None,
    tokenizer_kwargs: Optional[Dict[str, Any]] = None,
    transformer_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
    super().__init__()
    from allennlp.common import cached_transformers

    self.transformer_model = cached_transformers.get(
        model_name,
        True,
        override_weights_file=override_weights_file,
        override_weights_strip_prefix=override_weights_strip_prefix,
        **(transformer_kwargs or {}),
    )

    if gradient_checkpointing is not None:
        self.transformer_model.config.update(
            {"gradient_checkpointing": gradient_checkpointing}
        )

    self.config = self.transformer_model.config
    if sub_module:
        assert hasattr(self.transformer_model, sub_module)
        self.transformer_model = getattr(self.transformer_model, sub_module)
    self._max_length = max_length

    # I'm not sure if this works for all models; open an issue on github if you find a case
    # where it doesn't work.
    self.output_dim = self.config.hidden_size

    self._scalar_mix: Optional[ScalarMix] = None
    if not last_layer_only:
        self._scalar_mix = ScalarMix(self.config.num_hidden_layers)
        self.config.output_hidden_states = True

    tokenizer = PretrainedTransformerTokenizer(
        model_name,
        tokenizer_kwargs=tokenizer_kwargs,
    )
    self._num_added_start_tokens = len(tokenizer.single_sequence_start_tokens)
    self._num_added_end_tokens = len(tokenizer.single_sequence_end_tokens)
    self._num_added_tokens = self._num_added_start_tokens + self._num_added_end_tokens

    if not train_parameters:
        for param in self.transformer_model.parameters():
            param.requires_grad = False


# Traceback (OSError)
2020-09-30 23:56:17,398 - INFO - allennlp.training.trainer - Epoch 0/9
2020-09-30 23:56:17,398 - INFO - allennlp.training.trainer - Worker 0 memory usage MB: 10065.304
2020-09-30 23:56:17,484 - WARNING - allennlp.common.util - unable to check gpu_memory_mb() due to occasional failure, continuing
Traceback (most recent call last):
File "/media/discoD/repositorios/allennlp/allennlp/common/util.py", line 415, in gpu_memory_mb
encoding="utf-8",
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/subprocess.py", line 411, in check_output
**kwargs).stdout
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/subprocess.py", line 488, in run
with Popen(*popenargs, **kwargs) as process:
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/subprocess.py", line 800, in __init__
restore_signals, start_new_session)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/subprocess.py", line 1482, in _execute_child
restore_signals, start_new_session, preexec_fn)
File "/media/discoD/pycharm-community-2019.2/plugins/python-ce/helpers/pydev/_pydev_bundle/pydev_monkey.py", line 526, in new_fork_exec
return getattr(_posixsubprocess, original_name)(args, *patch_fork_exec_executable_list(args, other_args))
OSError: [Errno 12] Cannot allocate memory
2020-09-30 23:56:17,489 - INFO - allennlp.training.trainer - Training
0%|          | 0/11817 [00:00<?, ?it/s]/pytorch/aten/src/THC/THCTensorIndex.cu:272: indexSelectLargeIndex: block: [69,0,0], thread: [32,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/pytorch/aten/src/THC/THCTensorIndex.cu:272: indexSelectLargeIndex: block: [69,0,0], thread: [33,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/pytorch/aten/src/THC/THCTensorIndex.cu:272: indexSelectLargeIndex: block: [69,0,0], thread: [34,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/pytorch/aten/src/THC/THCTensorIndex.cu:272: indexSelectLargeIndex: block: [69,0,0], thread: [35,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
...
...
...
/pytorch/aten/src/THC/THCTensorIndex.cu:272: indexSelectLargeIndex: block: [102,0,0], thread: [30,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
/pytorch/aten/src/THC/THCTensorIndex.cu:272: indexSelectLargeIndex: block: [102,0,0], thread: [31,0,0] Assertion `srcIndex < srcSelectDimSize` failed.
0%|          | 0/11817 [00:00<?, ?it/s]
Traceback (most recent call last):
File "/media/discoD/repositorios/allennlp/allennlp/commands/train.py", line 443, in _train_worker
metrics = train_loop.run()
File "/media/discoD/repositorios/allennlp/allennlp/commands/train.py", line 505, in run
return self.trainer.train()
File "/media/discoD/repositorios/allennlp/allennlp/training/trainer.py", line 872, in train
train_metrics = self._train_epoch(epoch)
File "/media/discoD/repositorios/allennlp/allennlp/training/trainer.py", line 594, in _train_epoch
batch_outputs = self.batch_outputs(batch, for_training=True)
File "/media/discoD/repositorios/allennlp/allennlp/training/trainer.py", line 479, in batch_outputs
output_dict = self._pytorch_model(**batch)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/repositorios/allennlp/allennlp/models/basic_classifier.py", line 121, in forward
embedded_text = self._text_field_embedder(tokens)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/repositorios/allennlp/allennlp/modules/text_field_embedders/basic_text_field_embedder.py", line 88, in forward
token_vectors = embedder(**tensors, **forward_params_values)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/repositorios/allennlp/allennlp/modules/token_embedders/pretrained_transformer_embedder.py", line 184, in forward
transformer_output = self.transformer_model(**parameters)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/transformers/modeling_bert.py", line 762, in forward
output_hidden_states=output_hidden_states,
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/transformers/modeling_bert.py", line 439, in forward
output_attentions,
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/transformers/modeling_bert.py", line 371, in forward
hidden_states, attention_mask, head_mask, output_attentions=output_attentions,
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/transformers/modeling_bert.py", line 315, in forward
hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, output_attentions,
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/transformers/modeling_bert.py", line 221, in forward
mixed_query_layer = self.query(hidden_states)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/module.py", line 722, in _call_impl
result = self.forward(*input, **kwargs)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/modules/linear.py", line 91, in forward
return F.linear(input, self.weight, self.bias)
File "/media/discoD/anaconda3/envs/allennlp/lib/python3.7/site-packages/torch/nn/functional.py", line 1676, in linear
output = input.matmul(weight.t())
RuntimeError: CUDA error: CUBLAS_STATUS_ALLOC_FAILED when calling `cublasCreate(handle)`
python-BaseException
THCudaCheck FAIL file=/pytorch/aten/src/THC/THCCachingHostAllocator.cpp line=278 error=710 : device-side assert triggered

# Fixed code
def __init__(
    self,
    model_name: str,
    *,
    max_length: int = None,
    sub_module: str = None,
    train_parameters: bool = True,
    last_layer_only: bool = True,
    override_weights_file: Optional[str] = None,
    override_weights_strip_prefix: Optional[str] = None,
    gradient_checkpointing: Optional[bool] = None,
    tokenizer_kwargs: Optional[Dict[str, Any]] = None,
    transformer_kwargs: Optional[Dict[str, Any]] = None,
) -> None:
    super().__init__()
    from allennlp.common import cached_transformers

    self.transformer_model = cached_transformers.get(
        model_name,
        True,
        override_weights_file=override_weights_file,
        override_weights_strip_prefix=override_weights_strip_prefix,
        **(transformer_kwargs or {}),
    )

    if gradient_checkpointing is not None:
        self.transformer_model.config.update(
            {"gradient_checkpointing": gradient_checkpointing}
        )

    self.config = self.transformer_model.config
    if sub_module:
        assert hasattr(self.transformer_model, sub_module)
        self.transformer_model = getattr(self.transformer_model, sub_module)
    self._max_length = max_length

    # I'm not sure if this works for all models; open an issue on github if you find a case
    # where it doesn't work.
    self.output_dim = self.config.hidden_size

    self._scalar_mix: Optional[ScalarMix] = None
    if not last_layer_only:
        self._scalar_mix = ScalarMix(self.config.num_hidden_layers)
        self.config.output_hidden_states = True

    tokenizer = PretrainedTransformerTokenizer(
        model_name,
        tokenizer_kwargs=tokenizer_kwargs,
    )

    try:
        if self.transformer_model.get_input_embeddings().num_embeddings != len(
            tokenizer.tokenizer
        ):
            self.transformer_model.resize_token_embeddings(len(tokenizer.tokenizer))
    except NotImplementedError:
        # Can't resize for transformers models that don't implement base_model.get_input_embeddings()
        logger.warning(
            "Could not resize the token embedding matrix of the transformer model. "
            "This model does not support resizing."
        )

    self._num_added_start_tokens = len(tokenizer.single_sequence_start_tokens)
    self._num_added_end_tokens = len(tokenizer.single_sequence_end_tokens)
    self._num_added_tokens = self._num_added_start_tokens + self._num_added_end_tokens

    if not train_parameters:
        for param in self.transformer_model.parameters():
            param.requires_grad = False