For_sequence_classification (p.kaun.1.0.0~alpha1.doc.kaun.models.Kaun_models.Bert.For_sequence

Sourceval create : ?config:config -> num_labels:int -> unit -> Kaun.module_

val forward : 
  model:Kaun.module_ ->
  params:'a Kaun.params ->
  input_ids:(int32, Rune.int32_elt) Rune.t ->
  ?attention_mask:(int32, Rune.int32_elt) Rune.t ->
  ?token_type_ids:(int32, Rune.int32_elt) Rune.t ->
  ?labels:(int32, Rune.int32_elt) Rune.t ->
  training:bool ->
  unit ->
  (float, 'a) Rune.t * (float, 'a) Rune.t option

Returns (logits, loss) where logits has shape batch_size; num_labels