Do you mean that I can modify my code as follows?
        with autocast():
               embs = model(block[0].to(device))[head_invmap]
               embs = embs.gather(
                    1,
                    block_types.view(-1, 1, 1).expand(embs.shape[0], 1, embs.shape[2]),
              )[:, 0]
              loss = nsloss(
                    block[0].dstdata[dgl.NID][head_invmap].to(device),
                    embs,
                    tails.to(device),
             )
            scaler.scale(loss).backward()
            scaler.step(optimizer)
            # although the paramaters don't contain any scaled paramaters. I still need to optimize it.
            sparse_optimizer.step()
            scaler.update()