mirrors/pytorch-image-models
1
0
Fork 0
mirror of https://github.com/huggingface/pytorch-image-models.git synced 2025-06-03 15:01:08 +08:00
Code Issues Projects Releases Wiki Activity

adapt_input_conv: add type hints

Browse Source
This commit is contained in:
Adam J. Stewart 2025-02-21 19:24:05 +01:00 committed by Ross Wightman
parent 105a667baa
commit c68d724e9c
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
timm/models/_manipulate.py
View File

@ -8,6 +8,7 @@ from typing import Any, Callable, Dict, Iterator, Optional, Tuple, Type, Union
import torch import torch
import torch.utils.checkpoint import torch.utils.checkpoint
from torch import nn as nn from torch import nn as nn
from torch import Tensor
from timm.layers import use_reentrant_ckpt from timm.layers import use_reentrant_ckpt
@ -284,7 +285,7 @@ def checkpoint_seq(
return x return x
def adapt_input_conv(in_chans, conv_weight): def adapt_input_conv(in_chans: int, conv_weight: Tensor) -> Tensor:
conv_type = conv_weight.dtype conv_type = conv_weight.dtype
conv_weight = conv_weight.float() # Some weights are in torch.half, ensure it's float for sum on CPU conv_weight = conv_weight.float() # Some weights are in torch.half, ensure it's float for sum on CPU
O, I, J, K = conv_weight.shape O, I, J, K = conv_weight.shape