bugfixes addressing issues with imports

makani/inference.py

@@ -53,6 +53,7 @@
  parser.add_argument("--epsilon_factor", default=0, type=float)
  parser.add_argument("--split_data_channels", action="store_true")
  parser.add_argument("--mode", default="score", type=str, choices=["score", "ensemble"], help="Select inference mode")
+ parser.add_argument("--enable_odirect", action="store_true")
 
  # checkpoint format
  parser.add_argument("--checkpoint_format", default="legacy", choices=["legacy", "flexible"], type=str, help="Format in which to load checkpoints.")
@@ -124,6 +125,7 @@
  params["amp_mode"] = args.amp_mode
  params["jit_mode"] = args.jit_mode
  params["cuda_graph_mode"] = args.cuda_graph_mode
+ params["enable_odirect"] = args.enable_odirect
  params["enable_benchy"] = args.enable_benchy
  params["disable_ddp"] = args.disable_ddp
  params["enable_nhwc"] = args.enable_nhwc

makani/utils/comm.py

@@ -93,7 +93,7 @@ def is_distributed(name: str):
  return False
 
 
-# initialization routine 
+# initialization routine
 def init(model_parallel_sizes=[1, 1, 1, 1],
  model_parallel_names=["h", "w", "fin", "fout"],
  verbose=False):

makani/utils/dataloader.py

@@ -72,7 +72,7 @@ def get_dataloader(params, files_pattern, device, train=True, final_eval=False):
  from makani.utils.dataloaders.data_loader_multifiles import MultifilesDataset as MultifilesDataset2D
  from torch.utils.data.distributed import DistributedSampler
 
- # multifiles dataset
+ # multifiles
  dataset = MultifilesDataset2D(params, files_pattern, train)
 
  sampler = DistributedSampler(dataset, shuffle=train, num_replicas=params.data_num_shards, rank=params.data_shard_id) if (params.data_num_shards > 1) else None
@@ -81,8 +81,8 @@ def get_dataloader(params, files_pattern, device, train=True, final_eval=False):
  dataset,
  batch_size=int(params.batch_size),
  num_workers=params.num_data_workers,
- shuffle=False, # (sampler is None),
- sampler=sampler if train else None,
+ shuffle=(sampler is None) and train,
+ sampler=sampler,
  drop_last=True,
  pin_memory=torch.cuda.is_available(),
  )

makani/utils/visualize.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
 import io
 import numpy as np
 import concurrent.futures as cf