v0.0.4 fn (%x: Tensor[(2, 4096), int32]) -> Tensor[(2, 4096), int32] { let %x1: Tensor[(2, 4096), int32] = nn.batch_flatten(%x); %x1 }