lama

1
from enum import Enum
2

3
import yaml
4
from easydict import EasyDict as edict
5
import torch.nn as nn
6
import torch
7

8

9
def load_yaml(path):
10
    with open(path, 'r') as f:
11
        return edict(yaml.safe_load(f))
12

13

14
def move_to_device(obj, device):
15
    if isinstance(obj, nn.Module):
16
        return obj.to(device)
17
    if torch.is_tensor(obj):
18
        return obj.to(device)
19
    if isinstance(obj, (tuple, list)):
20
        return [move_to_device(el, device) for el in obj]
21
    if isinstance(obj, dict):
22
        return {name: move_to_device(val, device) for name, val in obj.items()}
23
    raise ValueError(f'Unexpected type {type(obj)}')
24

25

26
class SmallMode(Enum):
27
    DROP = "drop"
28
    UPSCALE = "upscale"
29