//  [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()

        B, N, _ = inputs.shape
        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)).reshape(B,-1)
        
        seq_length = inputs.size()[1]
        position_ids_start = torch.arange(seq_length, dtype=torch.long, device=inputs.device).view(1, -1)
        position_ids_end = torch.arange(seq_length, dtype=torch.long, device=inputs.device).view(1, -1)
        
        position_ids_start = (position_ids_start - start_idx[:,None])+seq_length
        position_ids_end = (position_ids_end - end_idx[:,None])+seq_length
        
        start_embedding = self.start_pos_embedding(position_ids_start)
        end_embedding = self.end_pos_embedding(position_ids_end)

        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).reshape(B, N,-1)
        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 + length_embed[:,None]+num_embed[:,None]+start_embedding+end_embedding)* 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
        )
        first_vector = True
        if first_vector:
            vectors = self.vector_decoder(sequence_output)
            B,N,_ = vectors.shape
            vectors = vectors.reshape(B,-1,3)
            vectors = vectors/torch.linalg.norm(vectors, dim=-1, keepdim=True)
            
            C_Ni = vectors[:,0::3]
            Ni_CAi = vectors[:,1::3]
            CAi_Ci = vectors[:,2::3] # CA80_C80*1.54 + C80_N81*1.34 + N81_CA81*1.46
            
            phi = self.dihedral(C_Ni, Ni_CAi, CAi_Ci) # OK
            psi = self.dihedral(Ni_CAi, CAi_Ci, torch.roll(C_Ni, -1, dims=1)) # OK
            omega = self.dihedral(CAi_Ci, torch.roll(C_Ni, -1, dims=1), torch.roll(Ni_CAi, -1, dims=1)) # OK
            tau = self.angle(-torch.roll(Ni_CAi, -1, dims=1), torch.roll(CAi_Ci, -1, dims=1)) # OK
            CA_C_1N = self.angle(-CAi_Ci, torch.roll(C_Ni,-1,dims=1)) # OK
            C_1N_1CA = self.angle(-torch.roll(C_Ni,-1,dims=1), torch.roll(Ni_CAi,-1,dims=1)) # OK
            
            pred = torch.stack([phi, psi, omega, tau, CA_C_1N, C_1N_1CA], dim=-1)
            
            N_N = self.get_vector(C_Ni, Ni_CAi, CAi_Ci, "N", "N", start_idx, end_idx)
            CA_CA = self.get_vector(C_Ni, Ni_CAi, CAi_Ci, "CA", "CA", start_idx, end_idx)
            C_C = self.get_vector(C_Ni, Ni_CAi, CAi_Ci, "C", "C", start_idx, end_idx)
            vectors = torch.stack([N_N, C_C, CA_CA], dim=1)
            
            # overlapping loss
            rand_idx = torch.cat([torch.randint(start_idx[i], end_idx[i],(1,)) for i in range(start_idx.shape[0])]).to(pred.device)
            
            anchor = coords[idx, start_idx,1] + self.get_vector(C_Ni, Ni_CAi, CAi_Ci, "CA", "CA", start_idx, rand_idx)
            
            dist = (anchor[:,None,None] - coords[:,None,:,1,:]).norm(dim=-1)[:,0]
            rel_idx = rand_idx[:,None] - torch.arange(coords.shape[1], device=coords.device)[None]
            select_mask = (~unknown_mask[:,:,0]* attention_mask*(rel_idx>10))
            dist = dist*select_mask + 10000*(1-select_mask)
            dist = dist.min(dim=1)[0]
        else:
            pred = self.token_decoder(sequence_output)
            phi, psi, omega, tau, CA_C_1N, C_1N_1CA = pred.split(1,dim=-1)
            coords_pred = self.pred_coord(coords, start_idx-1, end_idx+1, phi, psi, omega, tau, CA_C_1N, C_1N_1CA)
            idx = torch.arange(start_idx.shape[0], device=coords.device)
            vectors = coords_pred[idx, end_idx] - coords_pred[idx, start_idx]
            
            rand_idx = torch.cat([torch.randint(start_idx[i], end_idx[i],(1,)) for i in range(start_idx.shape[0])]).to(pred.device)
            
            # print(coords[idx,rand_idx,1].shape,  coords[:,None,:,1,:].shape)
            dist = (coords[idx,rand_idx,1][:,None,None] - coords[:,None,:,1,:]).norm(dim=-1)[:,0]
            rel_idx = rand_idx[:,None] - torch.arange(coords.shape[1], device=coords.device)[None]
            select_mask = (~unknown_mask[:,:,0]* attention_mask*(rel_idx>10))
            dist = dist*select_mask + 10000*(1-select_mask)
            dist = dist.min(dim=1)[0]
            
        return pred, vectors, dist