//  [python]
def forward(
        self,
        inputs: torch.Tensor,
        coords,
        timestep: torch.Tensor,  # Tensor of shape batch_length with time indices
        attention_mask: torch.Tensor,
        position_ids: Optional[torch.Tensor],
        seqs: torch.Tensor, 
        unknown_mask: torch.Tensor, 
        start_idx, 
        end_idx,
        head_mask: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        
        t1 = time.time()


        num = end_idx - start_idx
        num_embed = self.num_embedding(num)
        idx = torch.arange(start_idx.shape[0], device = start_idx.device)
        length = torch.linalg.norm(coords[idx, end_idx,1,:] - coords[idx, start_idx,1,:], dim=1)
        # length_embed = self.length_embedding(_rbf(length, 32).squeeze())

        output_attentions = (
            output_attentions
            if output_attentions is not None
            else self.config.output_attentions
        )
        output_hidden_states = (
            output_hidden_states
            if output_hidden_states is not None
            else self.config.output_hidden_states
        )
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        input_shape = inputs.size()
        batch_size, seq_length, *_ = input_shape
        logging.debug(f"Detected batch {batch_size} and seq length {seq_length}")

        assert attention_mask is not None

        # If position IDs are not given, auto-generate them
        if position_ids is None:
            # [1, seq_length]
            position_ids = (
                torch.arange(
                    seq_length,
                )
                .expand(batch_size, -1)
                .type_as(timestep)
            )


        assert (
            attention_mask.dim() == 2
        ), f"Attention mask expected in shape (batch_size, seq_length), got {attention_mask.shape}"
        extended_attention_mask = attention_mask[:, None, None, :]
        extended_attention_mask = extended_attention_mask.type_as(attention_mask)
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0


        assert len(inputs.shape) == 3  # batch_size, seq_length, features
        inputs_upscaled = self.inputs_to_hidden_dim(inputs)  # Batch * seq_len * dim

        # Pass through embeddings
        inputs_upscaled = self.embeddings(inputs_upscaled, position_ids=position_ids)
        seq_encoded = self.seq_embedding(seqs).squeeze()
        mask_encoded = self.mask_embedding(unknown_mask.long()).squeeze()

        # timestep is (batch, 1), squeeze to (batch,)
        # embedding gets to (batch, embed_dim) -> unsqueee to (batch, 1, dim)
        time_encoded = self.time_embed(timestep.squeeze(dim=-1)).unsqueeze(1) # [64, 1, 384]
        inputs_with_time = (inputs_upscaled + time_encoded + seq_encoded)* torch.sigmoid(mask_encoded)

        sequence_output = self.encoder(
            inputs_with_time,
            attention_mask=extended_attention_mask,
            # head_mask=head_mask,
            output_attentions=True,
            return_dict=return_dict,
            att_bias = None
        )

        noise = self.token_decoder(sequence_output)
        
        pred = utils.modulo_with_wrapped_range(inputs + noise, -torch.pi, torch.pi)

        t2 = time.time()
        # print(t2-t1)
        return pred