diff --git a/codebert.py b/codebert.py index db11eac..de3ed5b 100644 --- a/codebert.py +++ b/codebert.py @@ -45,7 +45,6 @@ class CodeBert(NNOperator): def __init__(self, model_name: str = 'huggingface/CodeBERTa-small-v1', device: str = None): super().__init__() self.model_name = model_name - assert modality in ['nlp', 'code'], 'Invalid modality value. Accept only "nlp" or "code".' if device is None: device = 'cuda' if torch.cuda.is_available() else 'cpu'