gr00t1.5_starforce / tests /test_tensor.py
nnh-pbbb's picture
Add files using upload-large-folder tool
cd793b5 verified
'''
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}")