'''
write a code read model weights from
checkpoints/models/pi0-sl-b1/model.safetensors
print all key and values
'''

'''
write a code read model weights from
checkpoints/models/pi0-sl-b1/model.safetensors
print all key and values
'''

from safetensors.torch import load_file

# Load the weights from the safetensors file
weights = load_file("checkpoints/models/pi0-sl-b1/model.safetensors")

# Print all keys and their corresponding values' shapes
for key, value in weights.items():
    if 'normalize' in key:
        print(f"Key: {key}, Shape: {value.shape}, Type: {value.dtype} {value}")