Merge branch 'KoboldAI:main' into united

This commit is contained in:
henk717 2022-07-12 23:03:09 +02:00 committed by GitHub
commit 68110c5930
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 3 additions and 2 deletions

View File

@ -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

View File

@ -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