File size: 323 Bytes
cb0ad2d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 |
from .model import Model
import torch
from torch import nn
from libs.utils.comm import get_world_size
def build_model(cfg):
if get_world_size() == 1:
norm_layer = nn.BatchNorm2d
else:
norm_layer = nn.BatchNorm2d
model = Model(
cfg,
norm_layer=norm_layer
)
return model
|