|
|
from typing import Dict, List, Tuple, Callable |
|
|
import torch |
|
|
import torch.nn as nn |
|
|
|
|
|
def get_module_device(m: nn.Module): |
|
|
device = torch.device('cpu') |
|
|
try: |
|
|
param = next(iter(m.parameters())) |
|
|
device = param.device |
|
|
except StopIteration: |
|
|
pass |
|
|
return device |
|
|
|
|
|
@torch.no_grad() |
|
|
def get_output_shape( |
|
|
input_shape: Tuple[int], |
|
|
net: Callable[[torch.Tensor], torch.Tensor] |
|
|
): |
|
|
device = get_module_device(net) |
|
|
test_input = torch.zeros((1,)+tuple(input_shape), device=device) |
|
|
test_output = net(test_input) |
|
|
output_shape = tuple(test_output.shape[1:]) |
|
|
return output_shape |
|
|
|
