Update utils.py

SongGeneration/codeclm/utils/utils.py

@@ -1,18 +1,8 @@
-from concurrent.futures import ProcessPoolExecutor
-from contextlib import contextmanager
-from functools import wraps, lru_cache
-import hashlib
-import json
-import logging
-from pathlib import Path
 import typing as tp
 import math
 from torch import nn
 import typing as tp
 from functools import partial
-import torch.nn.functional as F
-import flashy 
-import flashy.distrib
 import omegaconf
 import torch
 from torch.nn.utils.rnn import pad_sequence
@@ -193,4 +183,4 @@ def multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, gen
     input_ = input.reshape(-1, input.shape[-1])
     output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator)
     output = output_.reshape(*list(input.shape[:-1]), -1)
-    return output
+    return output