Retrieve more information from PyTorch checkpoints. (#515)

* Retrieve more information from PyTorch checkpoints.

* Add enough support to load dino-v2 backbone weights.
This commit is contained in:
Laurent Mazare
2023-08-19 15:05:34 +01:00
committed by GitHub
parent f861a9df6e
commit 607ffb9f1e
3 changed files with 75 additions and 20 deletions

View File

@ -9,6 +9,14 @@ pub struct Layout {
}
impl Layout {
pub fn new(shape: Shape, stride: Vec<usize>, start_offset: usize) -> Self {
Self {
shape,
stride,
start_offset,
}
}
pub fn contiguous_with_offset<S: Into<Shape>>(shape: S, start_offset: usize) -> Self {
let shape = shape.into();
let stride = shape.stride_contiguous();