Support importing safetensors format

Cargo.toml

@@ -8,8 +8,8 @@ members = [
     "crates/*",
     "crates/burn-import/pytorch-tests",
     "crates/burn-import/onnx-tests",
+    "crates/burn-import/safetensors-tests",
     "examples/*",
-    "examples/pytorch-import/model",
     "xtask",
 ]
 

README.md

@@ -188,18 +188,18 @@ for more details.
 <div align="left">
 <img align="right" src="https://raw.github.com/tracel-ai/burn/main/assets/backend-chip.png" height="96px"/>
 
-
 Burn strives to be as fast as possible on as many hardwares as possible, with robust implementations.
 We believe this flexibility is crucial for modern needs where you may train your models in the cloud,
 then deploy on customer hardwares, which vary from user to user.
+
 </div>
 
 <br />
 
 **Supported Backends**
 
 | Backend  | Devices                      | Class       |
-| -------  | ---------------------------- | ----------- |
+| -------- | ---------------------------- | ----------- |
 | CUDA     | NVIDIA GPUs                  | First-Party |
 | ROCm     | AMD GPUs                     | First-Party |
 | Metal    | Apple GPUs                   | First-Party |
@@ -426,13 +426,16 @@ Our ONNX support is further described in
 
 <details>
 <summary>
-Importing PyTorch Models 🚚
+Importing PyTorch or Safetensors Models 🚚
 </summary>
 <br />
 
-Support for loading of PyTorch model weights into Burn’s native model architecture, ensuring
-seamless integration. See
-[Burn Book 🔥 section on importing PyTorch](https://burn.dev/burn-book/import/pytorch-model.html)
+You can load weights from PyTorch or Safetensors formats directly into your Burn-defined models. This makes it easy to reuse existing models while benefiting from Burn's performance and deployment features.
+
+Learn more:
+
+- [Import pre-trained PyTorch models into Burn](https://burn.dev/burn-book/import/pytorch-model.html)
+- [Load models from Safetensors format](https://burn.dev/burn-book/import/safetensors-model.html)
 
 </details>
 
@@ -468,24 +471,23 @@ means it can run in bare metal environment such as embedded devices without an o
 
 <br />
 
-
 ### Benchmarks
 
 To evaluate performance across different backends and track improvements over time, we provide a
 dedicated benchmarking suite.
 
 Run and compare benchmarks using [burn-bench](https://github.com/tracel-ai/burn-bench).
 
-
-> ⚠️ **Warning**  
-> When using one of the `wgpu` backends, you may encounter compilation errors related to recursive type evaluation. This is due to complex type nesting within the `wgpu` dependency chain.  
+> ⚠️ **Warning**
+> When using one of the `wgpu` backends, you may encounter compilation errors related to recursive type evaluation. This is due to complex type nesting within the `wgpu` dependency chain.
 > To resolve this issue, add the following line at the top of your `main.rs` or `lib.rs` file:
+>
 > ```rust
 > #![recursion_limit = "256"]
 > ```
+>
 > The default recursion limit (128) is often just below the required depth (typically 130-150) due to deeply nested associated types and trait bounds.
 
-
 ## Getting Started
 
 <div align="left">

burn-book/src/SUMMARY.md

@@ -24,6 +24,7 @@
 - [Import Models](./import/README.md)
   - [ONNX Model](./import/onnx-model.md)
   - [PyTorch Model](./import/pytorch-model.md)
+  - [Safetensors Model](./import/safetensors-model.md)
 - [Models & Pre-Trained Weights](./models-and-pretrained-weights.md)
 - [Quantization (Beta)](./quantization.md)
 - [Advanced](./advanced/README.md)

burn-book/src/import/README.md

@@ -1,10 +1,13 @@
 # Importing Models
 
-The Burn project supports the import of models from various frameworks, emphasizing efficiency and
-compatibility. Currently, it handles two primary model formats:
+Burn supports importing models from other frameworks and file formats, enabling you to use pre-trained weights in your Burn applications.
 
-1. [ONNX](./onnx-model.md): Facilitates direct import, ensuring the model's performance and structure
-   are maintained.
+## Supported Formats
 
-2. [PyTorch](./pytorch-model.md): Enables the loading of PyTorch model weights into Burn’s native model
-   architecture, ensuring seamless integration.
+Burn currently supports three primary model import formats:
+
+| Format | Description | Use Case |
+|--------|-------------|----------|
+| [**ONNX**](./onnx-model.md) | Open Neural Network Exchange format | Direct import of complete model architectures and weights from any framework that supports ONNX export |
+| [**PyTorch**](./pytorch-model.md) | PyTorch weights (.pt, .pth) | Loading weights from PyTorch models into a matching Burn architecture |
+| [**Safetensors**](./safetensors-model.md) | Hugging Face's model serialization format | Loading a model's tensor weights into a matching Burn architecture |