`torch_dtype` is deprecated! Use `dtype` instead!
You are attempting to use Flash Attention 2 without specifying a torch dtype. This might lead to unexpected behaviour
/root/eval/work/Isaac-GR00T/gr00t/model/modules/dit.py:255: FutureWarning: Accessing config attribute `compute_dtype` directly via 'AlternateVLDiT' object attribute is deprecated. Please access 'compute_dtype' over 'AlternateVLDiT's config object instead, e.g. 'unet.config.compute_dtype'.
  embedding_dim=self.inner_dim, compute_dtype=self.compute_dtype
/root/eval/work/Isaac-GR00T/gr00t/model/modules/dit.py:286: FutureWarning: Accessing config attribute `output_dim` directly via 'AlternateVLDiT' object attribute is deprecated. Please access 'output_dim' over 'AlternateVLDiT's config object instead, e.g. 'unet.config.output_dim'.
  self.proj_out_2 = nn.Linear(self.inner_dim, self.output_dim)
Total number of DiT parameters:  1091722240
Total number of SelfAttentionTransformer parameters:  201433088

Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]
Loading checkpoint shards:  33%|███▎      | 1/3 [00:00<00:00,  3.84it/s]
Loading checkpoint shards:  67%|██████▋   | 2/3 [00:00<00:00,  2.33it/s]
Loading checkpoint shards: 100%|██████████| 3/3 [00:01<00:00,  2.54it/s]
Loading checkpoint shards: 100%|██████████| 3/3 [00:01<00:00,  2.59it/s]
Complete checkpoint restoration verified: {"complete": true, "loading_info": {"missing_keys": [], "unexpected_keys": [], "mismatched_keys": [], "error_msgs": []}, "parameter_objects": 1030, "checkpoint_index_tensors": 1030, "backbone_config_sha256": "bec4b3d446efa05807365c9e1cec03ac590836879d02f3a6da879971154bdd3b"}
Serving GrootBackend on 127.0.0.1:5555