Failed to import adam2
--- Benchmarking on CUDA ---
--- Precision: FLOAT32 ---
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
--------------------------------------------------------------------------------
Size: 32x32, Batch: 32, Precision: FP32
Model                | Params       | Throughput (samples/s)    | Mode      
--------------------------------------------------------------------------------
Basic CNN            | 8.90 M       | 70425.93                  | FP32      
ResNet18             | 11.30 M      | 37859.00                  | FP32      
ResNet50             | 24.03 M      | 21963.30                  | FP32      
EfficientNet-B2      | 8.06 M       | 24561.31                  | FP32      
DiT-Medium           | 2.76 M       | 470.90                    | FP32      
TRM-Light            | 230 K        | 8245.65                   | Native    
TRM-Medium           | 1.44 M       | 720.83                    | Native    
TRM-Heavy            | 9.18 M       | 71.84                     | Native    
--------------------------------------------------------------------------------
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
--------------------------------------------------------------------------------
Size: 46x46, Batch: 32, Precision: FP32
Model                | Params       | Throughput (samples/s)    | Mode      
--------------------------------------------------------------------------------
Basic CNN            | 8.90 M       | 50495.76                  | FP32      
ResNet18             | 11.30 M      | 29884.60                  | FP32      
ResNet50             | 24.03 M      | 10985.96                  | FP32      
EfficientNet-B2      | 8.06 M       | 11040.91                  | FP32      
DiT-Medium           | 3.04 M       | 135.65                    | FP32      
TRM-Light            | 230 K        | 2279.37                   | Native    
TRM-Medium           | 1.44 M       | 195.13                    | Native    
TRM-Heavy            | 9.18 M       | 20.84                     | Native    
--------------------------------------------------------------------------------
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
--------------------------------------------------------------------------------
Size: 64x64, Batch: 32, Precision: FP32
Model                | Params       | Throughput (samples/s)    | Mode      
--------------------------------------------------------------------------------
Basic CNN            | 8.90 M       | 37015.37                  | FP32      
ResNet18             | 11.30 M      | 28849.14                  | FP32      
ResNet50             | 24.03 M      | 11313.02                  | FP32      
EfficientNet-B2      | 8.06 M       | 9409.81                   | FP32      
DiT-Medium           | 3.55 M       | 42.29                     | FP32      
TRM-Light            | 230 K        | 698.56                    | Native    
TRM-Medium           | 1.44 M       | 58.57                     | Native    
TRM-Heavy            | 9.18 M       | 6.56                      | Native    
--------------------------------------------------------------------------------
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
--------------------------------------------------------------------------------
Size: 90x90, Batch: 32, Precision: FP32
Model                | Params       | Throughput (samples/s)    | Mode      
--------------------------------------------------------------------------------
Basic CNN            | 8.90 M       | 24743.33                  | FP32      
ResNet18             | 11.30 M      | 23654.04                  | FP32      
ResNet50             | 24.03 M      | 7843.76                   | FP32      
EfficientNet-B2      | 8.06 M       | 7245.46                   | FP32      
DiT-Medium           | 4.57 M       | 11.40                     | FP32      
TRM-Light            | 230 K        | 186.14                    | Native    
TRM-Medium           | 1.44 M       | 15.55                     | Native    
TRM-Heavy            | 9.18 M       | 1.80                      | Native    
--------------------------------------------------------------------------------
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
Manually casting TRM to torch.float32
--------------------------------------------------------------------------------
Size: 128x128, Batch: 32, Precision: FP32
Model                | Params       | Throughput (samples/s)    | Mode      
--------------------------------------------------------------------------------
Basic CNN            | 8.90 M       | 17173.49                  | FP32      
ResNet18             | 11.30 M      | 22351.00                  | FP32      
ResNet50             | 24.03 M      | 8031.70                   | FP32      
EfficientNet-B2      | 8.06 M       | 6756.70                   | FP32      
DiT-Medium           | 6.69 M       | 2.91                      | FP32      
TRM-Light            | 230 K        | 47.03                     | Native    
TRM-Medium           | 1.44 M       | 3.91                      | Native    
TRM-Heavy            | 9.18 M       | 0.46                      | Native    
--------------------------------------------------------------------------------
