Searched refs:model_cls (Results 1 – 2 of 2) sorted by relevance
130 def get_sequence_length(model_cls, model_name): argument172 model_cls, model, model_name, bs, device, include_loss_args=False argument177 seq_length = get_sequence_length(model_cls, model_name)207 or model_cls233 if model_cls in [ElectraForPreTraining, LxmertForPreTraining]:238 if model_cls in [AlbertForPreTraining]369 model_cls = get_module_cls_by_model_name(model_name)370 config_cls = model_cls.config_class375 model_cls381 or model_cls.__name__.startswith("Roberta")[all …]
189 model_cls = NonUniformReqGradNWM193 model_cls = NestedWrappedModule # type: ignore[assignment]197 model_cls,