mirror of
https://github.com/KoboldAI/KoboldAI-Client.git
synced 2025-02-09 08:18:40 +01:00
Merge branch 'KoboldAI:main' into united
This commit is contained in:
commit
68110c5930
@ -5,6 +5,7 @@ requests
|
|||||||
optax >= 0.0.5, <= 0.0.9
|
optax >= 0.0.5, <= 0.0.9
|
||||||
dm-haiku == 0.0.5
|
dm-haiku == 0.0.5
|
||||||
jax == 0.2.21
|
jax == 0.2.21
|
||||||
|
jaxlib >= 0.1.69, <= 0.3.7
|
||||||
transformers >= 4.19
|
transformers >= 4.19
|
||||||
progressbar2
|
progressbar2
|
||||||
git+https://github.com/VE-FORBRYDERNE/mesh-transformer-jax@ck
|
git+https://github.com/VE-FORBRYDERNE/mesh-transformer-jax@ck
|
||||||
|
@ -52,7 +52,7 @@ import pickle
|
|||||||
import torch
|
import torch
|
||||||
import utils
|
import utils
|
||||||
from torch.nn import Module
|
from torch.nn import Module
|
||||||
from typing import Any, Callable, Dict, Optional, Tuple, Type, Union
|
from typing import Any, Callable, Dict, Optional, Tuple, Union
|
||||||
|
|
||||||
|
|
||||||
_EXTRA_STATE_KEY_SUFFIX = '_extra_state'
|
_EXTRA_STATE_KEY_SUFFIX = '_extra_state'
|
||||||
@ -73,7 +73,7 @@ STORAGE_TYPE_MAP = {
|
|||||||
|
|
||||||
|
|
||||||
class LazyTensor:
|
class LazyTensor:
|
||||||
def __init__(self, storage_type: Type[torch._StorageBase], key: str, location: str, dtype: Optional[torch.dtype] = None, seek_offset: Optional[int] = None, shape: Optional[Tuple[int, ...]] = None, stride: Optional[Tuple[int, ...]] = None, requires_grad=False, backward_hooks: Any = None):
|
def __init__(self, storage_type, key: str, location: str, dtype: Optional[torch.dtype] = None, seek_offset: Optional[int] = None, shape: Optional[Tuple[int, ...]] = None, stride: Optional[Tuple[int, ...]] = None, requires_grad=False, backward_hooks: Any = None):
|
||||||
self.storage_type = storage_type
|
self.storage_type = storage_type
|
||||||
self.key = key
|
self.key = key
|
||||||
self.location = location
|
self.location = location
|
||||||
|
Loading…
x
Reference in New Issue
Block a user